<commit_before>package genapp

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"text/template"

	"github.com/goadesign/goa/design"
	"github.com/goadesign/goa/goagen/codegen"
)

func makeTestDir(g *Generator, apiName string) (outDir string, err error) {
	outDir = filepath.Join(g.outDir, "test")
	if err = os.RemoveAll(outDir); err != nil {
		return
	}
	if err = os.MkdirAll(outDir, 0755); err != nil {
		return
	}
	g.genfiles = append(g.genfiles, outDir)
	return
}

// TestMethod structure
type TestMethod struct {
	Name           string
	Comment        string
	ResourceName   string
	ActionName     string
	ControllerName string
	ContextVarName string
	ContextType    string
	RouteVerb      string
	FullPath       string
	Status         int
	ReturnType     *ObjectType
	Params         []ObjectType
	Payload        *ObjectType
}

// ObjectType structure
type ObjectType struct {
	Label       string
	Name        string
	Type        string
	Pointer     string
	Validatable bool
}

func (g *Generator) generateResourceTest(api *design.APIDefinition) error {
	if len(api.Resources) == 0 {
		return nil
	}
	testTmpl := template.Must(template.New("resources").Parse(testTmpl))
	outDir, err := makeTestDir(g, api.Name)
	if err != nil {
		return err
	}
	appPkg, err := codegen.PackagePath(g.outDir)
	if err != nil {
		return err
	}
	imports := []*codegen.ImportSpec{
		codegen.SimpleImport("bytes"),
		codegen.SimpleImport("fmt"),
		codegen.SimpleImport("io"),
		codegen.SimpleImport("log"),
		codegen.SimpleImport("net/http"),
		codegen.SimpleImport("net/http/httptest"),
		codegen.SimpleImport("net/url"),
		codegen.SimpleImport("testing"),
		codegen.SimpleImport(appPkg),
		codegen.SimpleImport("github.com/goadesign/goa"),
		codegen.SimpleImport("github.com/goadesign/goa/goatest"),
		codegen.SimpleImport("golang.org/x/net/context"),
	}

	return api.IterateResources(func(res *design.ResourceDefinition) error {
		filename := filepath.Join(outDir, codegen.SnakeCase(res.Name)+".go")
		file, err := codegen.SourceFileFor(filename)
		if err != nil {
			return err
		}
		if err := file.WriteHeader("", "test", imports); err != nil {
			return err
		}

		var methods []*TestMethod

		if err := res.IterateActions(func(action *design.ActionDefinition) error {
			if err := action.IterateResponses(func(response *design.ResponseDefinition) error {
				if response.Status == 101 { // SwitchingProtocols, Don't currently handle WebSocket endpoints
					return nil
				}
				for routeIndex, route := range action.Routes {
					mediaType := design.Design.MediaTypeWithIdentifier(response.MediaType)
					if mediaType == nil {
						methods = append(methods, g.createTestMethod(res, action, response, route, routeIndex, nil, nil))
					} else {
						if err := mediaType.IterateViews(func(view *design.ViewDefinition) error {
							methods = append(methods, g.createTestMethod(res, action, response, route, routeIndex, mediaType, view))
							return nil
						}); err != nil {
							return err
						}
					}
				}
				return nil
			}); err != nil {
				return err
			}
			return nil
		}); err != nil {
			return err
		}
		g.genfiles = append(g.genfiles, filename)
		err = testTmpl.Execute(file, methods)
		if err != nil {
			panic(err)
		}
		return file.FormatCode()
	})
}

func (g *Generator) createTestMethod(resource *design.ResourceDefinition, action *design.ActionDefinition,
	response *design.ResponseDefinition, route *design.RouteDefinition, routeIndex int,
	mediaType *design.MediaTypeDefinition, view *design.ViewDefinition) *TestMethod {

	var (
		actionName, ctrlName, varName                string
		routeQualifier, viewQualifier, respQualifier string
		comment                                      string
		returnType                                   *ObjectType
		params                                       []ObjectType
		payload                                      *ObjectType
	)

	actionName = codegen.Goify(action.Name, true)
	ctrlName = codegen.Goify(resource.Name, true)
	varName = codegen.Goify(action.Name, false)
	routeQualifier = suffixRoute(action.Routes, routeIndex)
	if view != nil && view.Name != "default" {
		viewQualifier = codegen.Goify(view.Name, true)
	}
	respQualifier = codegen.Goify(response.Name, true)
	hasReturnValue := view != nil && mediaType != nil

	if hasReturnValue {
		p, _, err := mediaType.Project(view.Name)
		if err != nil {
			panic(err) // bug
		}
		tmp := codegen.GoTypeName(p, nil, 0, false)
		if !p.IsBuiltIn() {
			tmp = fmt.Sprintf("%s.%s", g.target, tmp)
		}
		validate := codegen.RecursiveChecker(p.AttributeDefinition, false, false, false, "payload", "raw", 1, true)
		returnType = &ObjectType{}
		returnType.Type = tmp
		returnType.Validatable = validate != ""
	}

	comment = "runs the method " + actionName + " of the given controller with the given parameters"
	if action.Payload != nil {
		comment += " and payload"
	}
	comment += ".\n// It returns the response writer so it's possible to inspect the response headers"
	if hasReturnValue {
		comment += " and the media type struct written to the response"
	}
	comment += "."

	for _, paramName := range route.Params() {
		for name, att := range action.Params.Type.ToObject() {
			if name == paramName {
				param := ObjectType{}
				param.Label = name
				param.Name = codegen.Goify(name, false)
				param.Type = codegen.GoTypeRef(att.Type, nil, 0, false)
				if att.Type.IsPrimitive() && action.Params.IsPrimitivePointer(name) {
					param.Pointer = "*"
				}
				params = append(params, param)
			}
		}
	}

	if action.Payload != nil {
		payload = &ObjectType{}
		payload.Name = "payload"
		payload.Type = fmt.Sprintf("%s.%s", g.target, codegen.Goify(action.Payload.TypeName, true))
		if !action.Payload.IsPrimitive() && !action.Payload.IsArray() && !action.Payload.IsHash() {
			payload.Pointer = "*"
		}

		validate := codegen.RecursiveChecker(action.Payload.AttributeDefinition, false, false, false, "payload", "raw", 1, false)
		if validate != "" {
			payload.Validatable = true
		}
	}

	return &TestMethod{
		Name:           fmt.Sprintf("%s%s%s%s%s", actionName, ctrlName, respQualifier, routeQualifier, viewQualifier),
		ActionName:     actionName,
		ResourceName:   ctrlName,
		Comment:        fmt.Sprintf("%s %s", actionName, comment),
		Params:         params,
		Payload:        payload,
		ReturnType:     returnType,
		ControllerName: fmt.Sprintf("%s.%sController", g.target, ctrlName),
		ContextVarName: fmt.Sprintf("%sCtx", varName),
		ContextType:    fmt.Sprintf("%s.New%s%sContext", g.target, actionName, ctrlName),
		RouteVerb:      route.Verb,
		Status:         response.Status,
		FullPath:       goPathFormat(route.FullPath()),
	}
}

func goPathFormat(path string) string {
	return design.WildcardRegex.ReplaceAllLiteralString(path, "/%v")
}

func suffixRoute(routes []*design.RouteDefinition, currIndex int) string {
	if len(routes) > 1 && currIndex > 0 {
		return strconv.Itoa(currIndex)
	}
	return ""
}

var testTmpl = `
{{ range $test := . }}
// {{ $test.Name }} {{ $test.Comment }}
// If ctx is nil then context.Background() is used.
// If service is nil then a default service is created.
func {{ $test.Name }}(t *testing.T, ctx context.Context, service *goa.Service, ctrl {{ $test.ControllerName}}{{/*
*/}}{{ range $param := $test.Params }}, {{ $param.Name }} {{ $param.Pointer }}{{ $param.Type }}{{ end }}{{/*
*/}}{{ if $test.Payload }}, {{ $test.Payload.Name }} {{ $test.Payload.Pointer }}{{ $test.Payload.Type }}{{ end }}){{/*
*/}} (http.ResponseWriter{{ if $test.ReturnType }}, *{{ $test.ReturnType.Type }}{{ end }}) {
	// Setup service
	var (
		logBuf bytes.Buffer
		resp   interface{}

		respSetter goatest.ResponseSetterFunc = func(r interface{}) { resp = r }
	)
	if service == nil {
		service = goatest.Service(&logBuf, respSetter)
	} else {
		logger := log.New(&logBuf, "", log.Ltime)
		service.WithLogger(goa.NewLogger(logger))
		newEncoder := func(io.Writer) goa.Encoder { return  respSetter }
		service.Encoder = goa.NewHTTPEncoder() // Make sure the code ends up using this decoder
		service.Encoder.Register(newEncoder, "*/*")
	}
{{ if $test.Payload }}{{ if $test.Payload.Validatable }}
	// Validate payload
	err := {{ $test.Payload.Name }}.Validate()
	if err != nil {
		e, ok := err.(*goa.Error)
		if !ok {
			panic(err) // bug
		}
		if e.Status != {{ $test.Status }} {
			t.Errorf("unexpected payload validation error: %+v", e)
		}
		{{ if $test.ReturnType }}return nil, nil{{ else }}return nil{{ end }}
	}
{{ end }}{{ end }}
	// Setup request context
	rw := httptest.NewRecorder()
	req, err := http.NewRequest("{{ $test.RouteVerb }}", fmt.Sprintf({{ printf "%q" $test.FullPath }}{{ range $param := $test.Params }}, {{ $param.Name }}{{ end }}), nil)
	if err != nil {
		panic("invalid test " + err.Error()) // bug
	}
	prms := url.Values{}
	{{ range $param := $test.Params }}prms["{{ $param.Label }}"] = []string{fmt.Sprintf("%v",{{ $param.Name}})}
	{{ end }}
	if ctx == nil {
		ctx = context.Background()
	}
	goaCtx := goa.NewContext(goa.WithAction(ctx, "{{ $test.ResourceName }}Test"), rw, req, prms)
	{{ $test.ContextVarName }}, err := {{ $test.ContextType }}(goaCtx, service)
	if err != nil {
		panic("invalid test data " + err.Error()) // bug
	}
	{{ if $test.Payload }}{{ $test.ContextVarName }}.Payload = {{ $test.Payload.Name }}{{ end }}

	// Perform action
	err = ctrl.{{ $test.ActionName}}({{ $test.ContextVarName }})

	// Validate response
	if err != nil {
		t.Fatalf("controller returned %s, logs:\n%s", err, logBuf.String())
	}
	if rw.Code != {{ $test.Status }} {
		t.Errorf("invalid response status code: got %+v, expected {{ $test.Status }}", rw.Code)
	}
{{ if $test.ReturnType }}	var mt *{{ $test.ReturnType.Type }}
	if resp != nil {
		var ok bool
		mt, ok = resp.(*{{ $test.ReturnType.Type }})
		if !ok {
			t.Errorf("invalid response media: got %+v, expected instance of {{ $test.ReturnType.Type }}", resp)
		}
{{ if $test.ReturnType.Validatable }}		err = mt.Validate()
		if err != nil {
			t.Errorf("invalid response media type: %s", err)
		}
{{ end }}	}
{{ end }}
	// Return results
	return rw{{ if $test.ReturnType }}, mt{{ end }}
}
{{ end }}`
<commit_msg>Use appropriate types for each endpoints (#580)<commit_after>package genapp

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"text/template"

	"github.com/goadesign/goa/design"
	"github.com/goadesign/goa/goagen/codegen"
)

func makeTestDir(g *Generator, apiName string) (outDir string, err error) {
	outDir = filepath.Join(g.outDir, "test")
	if err = os.RemoveAll(outDir); err != nil {
		return
	}
	if err = os.MkdirAll(outDir, 0755); err != nil {
		return
	}
	g.genfiles = append(g.genfiles, outDir)
	return
}

// TestMethod structure
type TestMethod struct {
	Name           string
	Comment        string
	ResourceName   string
	ActionName     string
	ControllerName string
	ContextVarName string
	ContextType    string
	RouteVerb      string
	FullPath       string
	Status         int
	ReturnType     *ObjectType
	Params         []ObjectType
	Payload        *ObjectType
}

// ObjectType structure
type ObjectType struct {
	Label       string
	Name        string
	Type        string
	Pointer     string
	Validatable bool
}

func (g *Generator) generateResourceTest(api *design.APIDefinition) error {
	if len(api.Resources) == 0 {
		return nil
	}
	testTmpl := template.Must(template.New("resources").Parse(testTmpl))
	outDir, err := makeTestDir(g, api.Name)
	if err != nil {
		return err
	}
	appPkg, err := codegen.PackagePath(g.outDir)
	if err != nil {
		return err
	}
	imports := []*codegen.ImportSpec{
		codegen.SimpleImport("bytes"),
		codegen.SimpleImport("fmt"),
		codegen.SimpleImport("io"),
		codegen.SimpleImport("log"),
		codegen.SimpleImport("net/http"),
		codegen.SimpleImport("net/http/httptest"),
		codegen.SimpleImport("net/url"),
		codegen.SimpleImport("testing"),
		codegen.SimpleImport(appPkg),
		codegen.SimpleImport("github.com/goadesign/goa"),
		codegen.SimpleImport("github.com/goadesign/goa/goatest"),
		codegen.SimpleImport("golang.org/x/net/context"),
	}

	return api.IterateResources(func(res *design.ResourceDefinition) error {
		filename := filepath.Join(outDir, codegen.SnakeCase(res.Name)+".go")
		file, err := codegen.SourceFileFor(filename)
		if err != nil {
			return err
		}
		if err := file.WriteHeader("", "test", imports); err != nil {
			return err
		}

		var methods []*TestMethod

		if err := res.IterateActions(func(action *design.ActionDefinition) error {
			if err := action.IterateResponses(func(response *design.ResponseDefinition) error {
				if response.Status == 101 { // SwitchingProtocols, Don't currently handle WebSocket endpoints
					return nil
				}
				for routeIndex, route := range action.Routes {
					mediaType := design.Design.MediaTypeWithIdentifier(response.MediaType)
					if mediaType == nil {
						methods = append(methods, g.createTestMethod(res, action, response, route, routeIndex, nil, nil))
					} else {
						if err := mediaType.IterateViews(func(view *design.ViewDefinition) error {
							methods = append(methods, g.createTestMethod(res, action, response, route, routeIndex, mediaType, view))
							return nil
						}); err != nil {
							return err
						}
					}
				}
				return nil
			}); err != nil {
				return err
			}
			return nil
		}); err != nil {
			return err
		}
		g.genfiles = append(g.genfiles, filename)
		err = testTmpl.Execute(file, methods)
		if err != nil {
			panic(err)
		}
		return file.FormatCode()
	})
}

func (g *Generator) createTestMethod(resource *design.ResourceDefinition, action *design.ActionDefinition,
	response *design.ResponseDefinition, route *design.RouteDefinition, routeIndex int,
	mediaType *design.MediaTypeDefinition, view *design.ViewDefinition) *TestMethod {

	var (
		actionName, ctrlName, varName                string
		routeQualifier, viewQualifier, respQualifier string
		comment                                      string
		returnType                                   *ObjectType
		params                                       []ObjectType
		payload                                      *ObjectType
	)

	actionName = codegen.Goify(action.Name, true)
	ctrlName = codegen.Goify(resource.Name, true)
	varName = codegen.Goify(action.Name, false)
	routeQualifier = suffixRoute(action.Routes, routeIndex)
	if view != nil && view.Name != "default" {
		viewQualifier = codegen.Goify(view.Name, true)
	}
	respQualifier = codegen.Goify(response.Name, true)
	hasReturnValue := view != nil && mediaType != nil

	if hasReturnValue {
		p, _, err := mediaType.Project(view.Name)
		if err != nil {
			panic(err) // bug
		}
		tmp := codegen.GoTypeName(p, nil, 0, false)
		if !p.IsBuiltIn() {
			tmp = fmt.Sprintf("%s.%s", g.target, tmp)
		}
		validate := codegen.RecursiveChecker(p.AttributeDefinition, false, false, false, "payload", "raw", 1, true)
		returnType = &ObjectType{}
		returnType.Type = tmp
		if p.IsObject() {
			returnType.Pointer = "*"
		}
		returnType.Validatable = validate != ""
	}

	comment = "runs the method " + actionName + " of the given controller with the given parameters"
	if action.Payload != nil {
		comment += " and payload"
	}
	comment += ".\n// It returns the response writer so it's possible to inspect the response headers"
	if hasReturnValue {
		comment += " and the media type struct written to the response"
	}
	comment += "."

	for _, paramName := range route.Params() {
		for name, att := range action.Params.Type.ToObject() {
			if name == paramName {
				param := ObjectType{}
				param.Label = name
				param.Name = codegen.Goify(name, false)
				param.Type = codegen.GoTypeRef(att.Type, nil, 0, false)
				if att.Type.IsPrimitive() && action.Params.IsPrimitivePointer(name) {
					param.Pointer = "*"
				}
				params = append(params, param)
			}
		}
	}

	if action.Payload != nil {
		payload = &ObjectType{}
		payload.Name = "payload"
		payload.Type = fmt.Sprintf("%s.%s", g.target, codegen.Goify(action.Payload.TypeName, true))
		if !action.Payload.IsPrimitive() && !action.Payload.IsArray() && !action.Payload.IsHash() {
			payload.Pointer = "*"
		}

		validate := codegen.RecursiveChecker(action.Payload.AttributeDefinition, false, false, false, "payload", "raw", 1, false)
		if validate != "" {
			payload.Validatable = true
		}
	}

	return &TestMethod{
		Name:           fmt.Sprintf("%s%s%s%s%s", actionName, ctrlName, respQualifier, routeQualifier, viewQualifier),
		ActionName:     actionName,
		ResourceName:   ctrlName,
		Comment:        fmt.Sprintf("%s %s", actionName, comment),
		Params:         params,
		Payload:        payload,
		ReturnType:     returnType,
		ControllerName: fmt.Sprintf("%s.%sController", g.target, ctrlName),
		ContextVarName: fmt.Sprintf("%sCtx", varName),
		ContextType:    fmt.Sprintf("%s.New%s%sContext", g.target, actionName, ctrlName),
		RouteVerb:      route.Verb,
		Status:         response.Status,
		FullPath:       goPathFormat(route.FullPath()),
	}
}

func goPathFormat(path string) string {
	return design.WildcardRegex.ReplaceAllLiteralString(path, "/%v")
}

func suffixRoute(routes []*design.RouteDefinition, currIndex int) string {
	if len(routes) > 1 && currIndex > 0 {
		return strconv.Itoa(currIndex)
	}
	return ""
}

var testTmpl = `
{{ range $test := . }}
// {{ $test.Name }} {{ $test.Comment }}
// If ctx is nil then context.Background() is used.
// If service is nil then a default service is created.
func {{ $test.Name }}(t *testing.T, ctx context.Context, service *goa.Service, ctrl {{ $test.ControllerName}}{{/*
*/}}{{ range $param := $test.Params }}, {{ $param.Name }} {{ $param.Pointer }}{{ $param.Type }}{{ end }}{{/*
*/}}{{ if $test.Payload }}, {{ $test.Payload.Name }} {{ $test.Payload.Pointer }}{{ $test.Payload.Type }}{{ end }}){{/*
*/}} (http.ResponseWriter{{ if $test.ReturnType }}, {{ $test.ReturnType.Pointer }}{{ $test.ReturnType.Type }}{{ end }}) {
	// Setup service
	var (
		logBuf bytes.Buffer
		resp   interface{}

		respSetter goatest.ResponseSetterFunc = func(r interface{}) { resp = r }
	)
	if service == nil {
		service = goatest.Service(&logBuf, respSetter)
	} else {
		logger := log.New(&logBuf, "", log.Ltime)
		service.WithLogger(goa.NewLogger(logger))
		newEncoder := func(io.Writer) goa.Encoder { return  respSetter }
		service.Encoder = goa.NewHTTPEncoder() // Make sure the code ends up using this decoder
		service.Encoder.Register(newEncoder, "*/*")
	}
{{ if $test.Payload }}{{ if $test.Payload.Validatable }}
	// Validate payload
	err := {{ $test.Payload.Name }}.Validate()
	if err != nil {
		e, ok := err.(*goa.Error)
		if !ok {
			panic(err) // bug
		}
		if e.Status != {{ $test.Status }} {
			t.Errorf("unexpected payload validation error: %+v", e)
		}
		{{ if $test.ReturnType }}return nil, nil{{ else }}return nil{{ end }}
	}
{{ end }}{{ end }}
	// Setup request context
	rw := httptest.NewRecorder()
	req, err := http.NewRequest("{{ $test.RouteVerb }}", fmt.Sprintf({{ printf "%q" $test.FullPath }}{{ range $param := $test.Params }}, {{ $param.Name }}{{ end }}), nil)
	if err != nil {
		panic("invalid test " + err.Error()) // bug
	}
	prms := url.Values{}
	{{ range $param := $test.Params }}prms["{{ $param.Label }}"] = []string{fmt.Sprintf("%v",{{ $param.Name}})}
	{{ end }}
	if ctx == nil {
		ctx = context.Background()
	}
	goaCtx := goa.NewContext(goa.WithAction(ctx, "{{ $test.ResourceName }}Test"), rw, req, prms)
	{{ $test.ContextVarName }}, err := {{ $test.ContextType }}(goaCtx, service)
	if err != nil {
		panic("invalid test data " + err.Error()) // bug
	}
	{{ if $test.Payload }}{{ $test.ContextVarName }}.Payload = {{ $test.Payload.Name }}{{ end }}

	// Perform action
	err = ctrl.{{ $test.ActionName}}({{ $test.ContextVarName }})

	// Validate response
	if err != nil {
		t.Fatalf("controller returned %s, logs:\n%s", err, logBuf.String())
	}
	if rw.Code != {{ $test.Status }} {
		t.Errorf("invalid response status code: got %+v, expected {{ $test.Status }}", rw.Code)
	}
{{ if $test.ReturnType }}	var mt {{ $test.ReturnType.Pointer }}{{ $test.ReturnType.Type }}
	if resp != nil {
		var ok bool
		mt, ok = resp.({{ $test.ReturnType.Pointer }}{{ $test.ReturnType.Type }})
		if !ok {
			t.Errorf("invalid response media: got %+v, expected instance of {{ $test.ReturnType.Type }}", resp)
		}
{{ if $test.ReturnType.Validatable }}		err = mt.Validate()
		if err != nil {
			t.Errorf("invalid response media type: %s", err)
		}
{{ end }}	}
{{ end }}
	// Return results
	return rw{{ if $test.ReturnType }}, mt{{ end }}
}
{{ end }}`
<|endoftext|>


<commit_before>// Queries a Gold instance and downloads all positive images for each tests.
// Stores metadata in the meta.json output file.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/fileutil"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/go/export"
	"go.skia.org/infra/golden/go/search"
	"go.skia.org/infra/golden/go/types"
)

const (
	// Default query that will get all tests.
	DEFAULT_QUERY = "fdiffmax=-1&fref=false&frgbamax=255&frgbamin=0&head=true&include=false&limit=50&match=gamma_correct&match=name&metric=combined&neg=false&offset=0&pos=true&query=source_type%3Dgm&sort=desc&unt=false&nodiff=true"

	// Default URL for the Gold instance we want to query.
	GOLD_URL = "https://gold.skia.org"

	// Search endpoint of Gold.
	SEARCH_PATH = "/json/search"

	// Template for image file names. Generally the digest is filled in.
	IMG_FILENAME_TMPL = "%s.png"

	// Name of the output metameta data file.
	META_DATA_FILE = "meta.json"
)

var (
	baseURL    = flag.String("base_url", GOLD_URL, "Query URL to retrieve meta data.")
	imgBaseURL = flag.String("img_url", GOLD_URL, "Url to retrieve the images.")
	indexFile  = flag.String("index_file", "./"+META_DATA_FILE, "Path of the index file.")
	outputDir  = flag.String("output_dir", "", "Directory where images should be written. If empty no images will be written.")
)

func main() {
	common.Init()

	// We need at least the index file or an output directory.
	if *outputDir == "" && *indexFile == "" {
		sklog.Fatal("No index file or output directory specified.")
	}

	// If the index file is empty write it to output directory.
	useIndexPath := *indexFile
	if useIndexPath == "" {
		useIndexPath = filepath.Join(*outputDir, META_DATA_FILE)
	}

	// Set up the http client.
	client := &http.Client{}

	// load the test meta data from Gold.
	testRecords, err := loadMetaData(client, *baseURL, DEFAULT_QUERY, META_DATA_FILE)
	if err != nil {
		sklog.Fatalf("Error loading meta data: %s", err)
	}
	sklog.Infoln("Meta data loaded from Gold.")

	// Write the index file to disk.
	if err := export.WriteTestRecords(useIndexPath, testRecords); err != nil {
		sklog.Fatalf("Error writing output file '%s': %s", filepath.Join(*outputDir, META_DATA_FILE), err)
	}
	sklog.Infoln("Index file written to disk.")

	// If an output directory was given, download the images referenced in the index file.
	if *outputDir != "" {
		if err := downloadImages(*outputDir, *imgBaseURL, client, testRecords); err != nil {
			sklog.Fatalf("Error downloading images: %s", err)
		}
	}

	sklog.Infof("Success. Knowledge data written to %s", *outputDir)
}

// loadMetaData makes a query to a Gold instance and parses the JSON response.
// It then groups images/digests by tests and returns them.
func loadMetaData(client *http.Client, baseURL, query, metaDataFileName string) ([]*export.TestRecord, error) {
	url := strings.TrimRight(baseURL, "/") + SEARCH_PATH + "?" + query
	sklog.Infof("Requesting url: %s", url)
	resp, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer util.Close(resp.Body)

	searchResp := &search.NewSearchResponse{}
	if err := json.NewDecoder(resp.Body).Decode(searchResp); err != nil {
		return nil, err
	}
	sklog.Infof("Meta data decoded successfully.")

	// Group the results by test.
	retMap := map[string]*export.TestRecord{}
	for _, oneDigest := range searchResp.Digests {
		testNameVal := oneDigest.ParamSet[types.PRIMARY_KEY_FIELD]
		if len(testNameVal) == 0 {
			sklog.Errorf("Error: Digest '%s' has no primaryKey in paramset", oneDigest.Digest)
			continue
		}

		digestInfo := &export.DigestInfo{
			SRDigest: oneDigest,
			URL:      export.DigestUrl(baseURL, oneDigest.Digest),
		}

		testName := oneDigest.ParamSet[types.PRIMARY_KEY_FIELD][0]
		if found, ok := retMap[testName]; ok {
			found.Digests = append(found.Digests, digestInfo)
		} else {
			retMap[testName] = &export.TestRecord{
				TestName: testName,
				Digests:  []*export.DigestInfo{digestInfo},
			}
		}
	}

	// Put the records into an array and return them.
	ret := make([]*export.TestRecord, 0, len(retMap))
	for _, oneTestRec := range retMap {
		ret = append(ret, oneTestRec)
	}

	return ret, nil
}

// downloadImages downloads all images referenced in the meta data to disk.
// One directory is created for each test.
func downloadImages(baseDir, baseURL string, client *http.Client, testRecs []*export.TestRecord) error {
	for _, testRec := range testRecs {
		testDir := filepath.Join(baseDir, testRec.TestName)
		absDirPath, err := fileutil.EnsureDirExists(testDir)
		if err != nil {
			sklog.Errorf("Error creating directory '%s'. Skipping. Got error: %s", testDir, err)
			continue
		}

		for _, digestRec := range testRec.Digests {
			// Download the image and then write it to disk.
			resp, err := client.Get(digestRec.URL)
			if err != nil {
				sklog.Errorf("Error retrieving file '%s'. Got error: %s", digestRec.URL, err)
				continue
			}
			defer util.Close(resp.Body)

			// Write the image to disk.
			imgFileName := fmt.Sprintf(IMG_FILENAME_TMPL, digestRec.Digest)
			fileName := filepath.Join(absDirPath, imgFileName)
			func(outputFileName string, reader io.Reader) {
				f, err := os.Create(outputFileName)
				if err != nil {
					sklog.Errorf("Error opening output file '%s'. Got error: %s", outputFileName, err)
					return
				}
				defer util.Close(f)

				if _, err := io.Copy(f, reader); err != nil {
					sklog.Errorf("Error saving file '%s'. Got error: %s", outputFileName, err)
					return
				}
				sklog.Infof("Downloaded %s sucessfully.", outputFileName)
			}(fileName, resp.Body)
		}
	}
	return nil
}
<commit_msg>Fix bug in knowledge export<commit_after>// Queries a Gold instance and downloads all positive images for each tests.
// Stores metadata in the meta.json output file.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/fileutil"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/go/export"
	"go.skia.org/infra/golden/go/search"
	"go.skia.org/infra/golden/go/types"
)

const (
	// Default query that will get all tests.
	DEFAULT_QUERY = "fdiffmax=-1&fref=false&frgbamax=255&frgbamin=0&head=true&include=false&limit=50&match=gamma_correct&match=name&metric=combined&neg=false&offset=0&pos=true&query=source_type%3Dgm&sort=desc&unt=false&nodiff=true"

	// Default URL for the Gold instance we want to query.
	GOLD_URL = "https://gold.skia.org"

	// Search endpoint of Gold.
	SEARCH_PATH = "/json/search"

	// Template for image file names. Generally the digest is filled in.
	IMG_FILENAME_TMPL = "%s.png"

	// Name of the output metameta data file.
	META_DATA_FILE = "meta.json"
)

var (
	baseURL    = flag.String("base_url", GOLD_URL, "Query URL to retrieve meta data.")
	imgBaseURL = flag.String("img_url", GOLD_URL, "Url to retrieve the images.")
	indexFile  = flag.String("index_file", "./"+META_DATA_FILE, "Path of the index file.")
	outputDir  = flag.String("output_dir", "", "Directory where images should be written. If empty no images will be written.")
)

func main() {
	common.Init()

	// We need at least the index file or an output directory.
	if *outputDir == "" && *indexFile == "" {
		sklog.Fatal("No index file or output directory specified.")
	}

	// If the index file is empty write it to output directory.
	useIndexPath := *indexFile
	if useIndexPath == "" {
		useIndexPath = filepath.Join(*outputDir, META_DATA_FILE)
	}

	// Set up the http client.
	client := &http.Client{}

	// load the test meta data from Gold.
	testRecords, err := loadMetaData(client, *baseURL, *imgBaseURL, DEFAULT_QUERY, META_DATA_FILE)
	if err != nil {
		sklog.Fatalf("Error loading meta data: %s", err)
	}
	sklog.Infoln("Meta data loaded from Gold.")

	// Write the index file to disk.
	if err := export.WriteTestRecords(useIndexPath, testRecords); err != nil {
		sklog.Fatalf("Error writing output file '%s': %s", filepath.Join(*outputDir, META_DATA_FILE), err)
	}
	sklog.Infoln("Index file written to disk.")

	// If an output directory was given, download the images referenced in the index file.
	if *outputDir != "" {
		if err := downloadImages(*outputDir, client, testRecords); err != nil {
			sklog.Fatalf("Error downloading images: %s", err)
		}
	}

	sklog.Infof("Success. Knowledge data written to %s", *outputDir)
}

// loadMetaData makes a query to a Gold instance and parses the JSON response.
// It then groups images/digests by tests and returns them.
func loadMetaData(client *http.Client, baseURL, imgBaseURL, query, metaDataFileName string) ([]*export.TestRecord, error) {
	url := strings.TrimRight(baseURL, "/") + SEARCH_PATH + "?" + query
	sklog.Infof("Requesting url: %s", url)
	resp, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer util.Close(resp.Body)

	searchResp := &search.NewSearchResponse{}
	if err := json.NewDecoder(resp.Body).Decode(searchResp); err != nil {
		return nil, err
	}
	sklog.Infof("Meta data decoded successfully.")

	// Group the results by test.
	retMap := map[string]*export.TestRecord{}
	for _, oneDigest := range searchResp.Digests {
		testNameVal := oneDigest.ParamSet[types.PRIMARY_KEY_FIELD]
		if len(testNameVal) == 0 {
			sklog.Errorf("Error: Digest '%s' has no primaryKey in paramset", oneDigest.Digest)
			continue
		}

		digestInfo := &export.DigestInfo{
			SRDigest: oneDigest,
			URL:      export.DigestUrl(imgBaseURL, oneDigest.Digest),
		}

		testName := oneDigest.ParamSet[types.PRIMARY_KEY_FIELD][0]
		if found, ok := retMap[testName]; ok {
			found.Digests = append(found.Digests, digestInfo)
		} else {
			retMap[testName] = &export.TestRecord{
				TestName: testName,
				Digests:  []*export.DigestInfo{digestInfo},
			}
		}
	}

	// Put the records into an array and return them.
	ret := make([]*export.TestRecord, 0, len(retMap))
	for _, oneTestRec := range retMap {
		ret = append(ret, oneTestRec)
	}

	return ret, nil
}

// downloadImages downloads all images referenced in the meta data to disk.
// One directory is created for each test.
func downloadImages(baseDir string, client *http.Client, testRecs []*export.TestRecord) error {
	for _, testRec := range testRecs {
		testDir := filepath.Join(baseDir, testRec.TestName)
		absDirPath, err := fileutil.EnsureDirExists(testDir)
		if err != nil {
			sklog.Errorf("Error creating directory '%s'. Skipping. Got error: %s", testDir, err)
			continue
		}

		for _, digestRec := range testRec.Digests {
			// Download the image and then write it to disk.
			resp, err := client.Get(digestRec.URL)
			if err != nil {
				sklog.Errorf("Error retrieving file '%s'. Got error: %s", digestRec.URL, err)
				continue
			}
			defer util.Close(resp.Body)

			// Write the image to disk.
			imgFileName := fmt.Sprintf(IMG_FILENAME_TMPL, digestRec.Digest)
			fileName := filepath.Join(absDirPath, imgFileName)
			func(outputFileName string, reader io.Reader) {
				f, err := os.Create(outputFileName)
				if err != nil {
					sklog.Errorf("Error opening output file '%s'. Got error: %s", outputFileName, err)
					return
				}
				defer util.Close(f)

				if _, err := io.Copy(f, reader); err != nil {
					sklog.Errorf("Error saving file '%s'. Got error: %s", outputFileName, err)
					return
				}
				sklog.Infof("Downloaded %s sucessfully.", outputFileName)
			}(fileName, resp.Body)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package nats_input

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/karimra/gnmic/formatters"
	"github.com/karimra/gnmic/inputs"
	"github.com/karimra/gnmic/outputs"
	"github.com/nats-io/nats.go"
	"google.golang.org/protobuf/proto"
)

const (
	loggingPrefix           = "nats_input "
	natsReconnectBufferSize = 100 * 1024 * 1024
	defaultAddress          = "localhost:4222"
	natsConnectWait         = 2 * time.Second
	defaultFormat           = "event"
	defaultSubject          = "telemetry"
	defaultNumWorkers       = 1
)

func init() {
	inputs.Register("nats", func() inputs.Input {
		return &NatsInput{
			Cfg:    &Config{},
			logger: log.New(ioutil.Discard, loggingPrefix, log.LstdFlags|log.Lmicroseconds),
			wg:     new(sync.WaitGroup),
		}
	})
}

// NatsInput //
type NatsInput struct {
	Cfg    *Config
	ctx    context.Context
	cfn    context.CancelFunc
	logger *log.Logger

	wg      *sync.WaitGroup
	outputs []outputs.Output
}

// Config //
type Config struct {
	Name            string        `mapstructure:"name,omitempty"`
	Address         string        `mapstructure:"address,omitempty"`
	Subject         string        `mapstructure:"subject,omitempty"`
	Queue           string        `mapstructure:"queue,omitempty"`
	Username        string        `mapstructure:"username,omitempty"`
	Password        string        `mapstructure:"password,omitempty"`
	ConnectTimeWait time.Duration `mapstructure:"connect-time-wait,omitempty"`
	Format          string        `mapstructure:"format,omitempty"`
	Debug           bool          `mapstructure:"debug,omitempty"`
	NumWorkers      int           `mapstructure:"num-workers,omitempty"`
	Outputs         []string      `mapstructure:"outputs,omitempty"`
}

// Init //
func (n *NatsInput) Start(ctx context.Context, cfg map[string]interface{}, opts ...inputs.Option) error {
	err := outputs.DecodeConfig(cfg, n.Cfg)
	if err != nil {
		return err
	}
	err = n.setDefaults()
	if err != nil {
		return err
	}
	for _, opt := range opts {
		opt(n)
	}
	n.ctx, n.cfn = context.WithCancel(ctx)
	n.logger.Printf("input starting with config: %+v", n.Cfg)
	n.wg.Add(n.Cfg.NumWorkers)
	for i := 0; i < n.Cfg.NumWorkers; i++ {
		go n.worker(ctx, i)
	}
	return nil
}

func (n *NatsInput) worker(ctx context.Context, idx int) {
	var nc *nats.Conn
	var err error
	var msgChan chan *nats.Msg
	workerLogPrefix := fmt.Sprintf("worker-%d", idx)
	n.logger.Printf("%s starting", workerLogPrefix)
	cfg := *n.Cfg
	cfg.Name = fmt.Sprintf("%s-%d", cfg.Name, idx)
START:
	nc, err = n.createNATSConn(&cfg)
	if err != nil {
		n.logger.Printf("%s failed to create NATS connection: %v", workerLogPrefix, err)
		time.Sleep(n.Cfg.ConnectTimeWait)
		goto START
	}
	defer nc.Close()
	msgChan = make(chan *nats.Msg)
	sub, err := nc.ChanSubscribe(n.Cfg.Subject, msgChan)
	if err != nil {
		n.logger.Printf("%s failed to create NATS connection: %v", workerLogPrefix, err)
		time.Sleep(n.Cfg.ConnectTimeWait)
		nc.Close()
		goto START
	}
	defer close(msgChan)
	defer sub.Unsubscribe()

	for {
		select {
		case <-ctx.Done():
			return
		case m, ok := <-msgChan:
			if !ok {
				n.logger.Printf("%s channel closed, retrying...", workerLogPrefix)
				time.Sleep(n.Cfg.ConnectTimeWait)
				nc.Close()
				goto START
			}
			if n.Cfg.Debug {
				n.logger.Printf("received msg, len=%d, data=%s", len(m.Data), string(m.Data))
			}

			switch n.Cfg.Format {
			case "event":
				evMsgs := make([]*formatters.EventMsg, 1)
				err = json.Unmarshal(m.Data, &evMsgs)
				if err != nil {
					if n.Cfg.Debug {
						n.logger.Printf("%s failed to unmarshal event msg: %v", workerLogPrefix, err)
					}
					continue
				}
				go func() {
					for _, o := range n.outputs {
						for _, ev := range evMsgs {
							o.WriteEvent(ctx, ev)
						}
					}
				}()
			case "proto":
				var protoMsg proto.Message
				err = proto.Unmarshal(m.Data, protoMsg)
				if err != nil {
					if n.Cfg.Debug {
						n.logger.Printf("failed to unmarshal proto msg: %v", err)
					}
					continue
				}
				meta := outputs.Meta{}
				subjectSections := strings.SplitN(m.Subject, ".", 3)
				if len(subjectSections) == 3 {
					meta["source"] = strings.ReplaceAll(subjectSections[1], "-", ".")
					meta["subscription-name"] = subjectSections[2]
				}
				go func() {
					for _, o := range n.outputs {
						o.Write(ctx, protoMsg, meta)
					}
				}()
			}

		}
	}
}

// Close //
func (n *NatsInput) Close() error {
	n.cfn()
	n.wg.Wait()
	return nil
}

// SetLogger //
func (n *NatsInput) SetLogger(logger *log.Logger) {
	if logger != nil && n.logger != nil {
		n.logger.SetOutput(logger.Writer())
		n.logger.SetFlags(logger.Flags())
	}
}

// SetOutputs //
func (n *NatsInput) SetOutputs(outs map[string]outputs.Output) {
	if len(n.Cfg.Outputs) == 0 {
		for _, o := range outs {
			n.outputs = append(n.outputs, o)
		}
		return
	}
	for _, name := range n.Cfg.Outputs {
		if o, ok := outs[name]; ok {
			n.outputs = append(n.outputs, o)
		}
	}
}

// helper functions

func (n *NatsInput) setDefaults() error {
	if n.Cfg.Format == "" {
		n.Cfg.Format = defaultFormat
	}
	if !(strings.ToLower(n.Cfg.Format) == "event" || strings.ToLower(n.Cfg.Format) == "proto") {
		return fmt.Errorf("unsupported input format")
	}
	if n.Cfg.Name == "" {
		n.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	if n.Cfg.Subject == "" {
		n.Cfg.Subject = defaultSubject
	}
	if n.Cfg.Address == "" {
		n.Cfg.Address = defaultAddress
	}
	if n.Cfg.ConnectTimeWait <= 0 {
		n.Cfg.ConnectTimeWait = natsConnectWait
	}
	if n.Cfg.Queue == "" {
		n.Cfg.Queue = n.Cfg.Name
	}
	if n.Cfg.NumWorkers <= 0 {
		n.Cfg.NumWorkers = defaultNumWorkers
	}
	return nil
}

func (n *NatsInput) createNATSConn(c *Config) (*nats.Conn, error) {
	opts := []nats.Option{
		nats.Name(c.Name),
		nats.SetCustomDialer(n),
		nats.ReconnectWait(n.Cfg.ConnectTimeWait),
		nats.ReconnectBufSize(natsReconnectBufferSize),
		nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) {
			n.logger.Printf("NATS error: %v", err)
		}),
		nats.DisconnectHandler(func(c *nats.Conn) {
			n.logger.Println("Disconnected from NATS")
		}),
		nats.ClosedHandler(func(c *nats.Conn) {
			n.logger.Println("NATS connection is closed")
		}),
	}
	if c.Username != "" && c.Password != "" {
		opts = append(opts, nats.UserInfo(c.Username, c.Password))
	}
	nc, err := nats.Connect(c.Address, opts...)
	if err != nil {
		return nil, err
	}
	return nc, nil
}

// Dial //
func (n *NatsInput) Dial(network, address string) (net.Conn, error) {
	ctx, cancel := context.WithCancel(n.ctx)
	defer cancel()

	for {
		n.logger.Printf("attempting to connect to %s", address)
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		select {
		case <-n.ctx.Done():
			return nil, n.ctx.Err()
		default:
			d := &net.Dialer{}
			if conn, err := d.DialContext(ctx, network, address); err == nil {
				n.logger.Printf("successfully connected to NATS server %s", address)
				return conn, nil
			}
			time.Sleep(n.Cfg.ConnectTimeWait)
		}
	}
}
<commit_msg>add queue subscription<commit_after>package nats_input

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/karimra/gnmic/formatters"
	"github.com/karimra/gnmic/inputs"
	"github.com/karimra/gnmic/outputs"
	"github.com/nats-io/nats.go"
	"google.golang.org/protobuf/proto"
)

const (
	loggingPrefix           = "nats_input "
	natsReconnectBufferSize = 100 * 1024 * 1024
	defaultAddress          = "localhost:4222"
	natsConnectWait         = 2 * time.Second
	defaultFormat           = "event"
	defaultSubject          = "telemetry"
	defaultNumWorkers       = 1
)

func init() {
	inputs.Register("nats", func() inputs.Input {
		return &NatsInput{
			Cfg:    &Config{},
			logger: log.New(ioutil.Discard, loggingPrefix, log.LstdFlags|log.Lmicroseconds),
			wg:     new(sync.WaitGroup),
		}
	})
}

// NatsInput //
type NatsInput struct {
	Cfg    *Config
	ctx    context.Context
	cfn    context.CancelFunc
	logger *log.Logger

	wg      *sync.WaitGroup
	outputs []outputs.Output
}

// Config //
type Config struct {
	Name            string        `mapstructure:"name,omitempty"`
	Address         string        `mapstructure:"address,omitempty"`
	Subject         string        `mapstructure:"subject,omitempty"`
	Queue           string        `mapstructure:"queue,omitempty"`
	Username        string        `mapstructure:"username,omitempty"`
	Password        string        `mapstructure:"password,omitempty"`
	ConnectTimeWait time.Duration `mapstructure:"connect-time-wait,omitempty"`
	Format          string        `mapstructure:"format,omitempty"`
	Debug           bool          `mapstructure:"debug,omitempty"`
	NumWorkers      int           `mapstructure:"num-workers,omitempty"`
	Outputs         []string      `mapstructure:"outputs,omitempty"`
}

// Init //
func (n *NatsInput) Start(ctx context.Context, cfg map[string]interface{}, opts ...inputs.Option) error {
	err := outputs.DecodeConfig(cfg, n.Cfg)
	if err != nil {
		return err
	}
	err = n.setDefaults()
	if err != nil {
		return err
	}
	for _, opt := range opts {
		opt(n)
	}
	n.ctx, n.cfn = context.WithCancel(ctx)
	n.logger.Printf("input starting with config: %+v", n.Cfg)
	n.wg.Add(n.Cfg.NumWorkers)
	for i := 0; i < n.Cfg.NumWorkers; i++ {
		go n.worker(ctx, i)
	}
	return nil
}

func (n *NatsInput) worker(ctx context.Context, idx int) {
	var nc *nats.Conn
	var err error
	var msgChan chan *nats.Msg
	workerLogPrefix := fmt.Sprintf("worker-%d", idx)
	n.logger.Printf("%s starting", workerLogPrefix)
	cfg := *n.Cfg
	cfg.Name = fmt.Sprintf("%s-%d", cfg.Name, idx)
START:
	nc, err = n.createNATSConn(&cfg)
	if err != nil {
		n.logger.Printf("%s failed to create NATS connection: %v", workerLogPrefix, err)
		time.Sleep(n.Cfg.ConnectTimeWait)
		goto START
	}
	defer nc.Close()
	msgChan = make(chan *nats.Msg)
	sub, err := nc.ChanQueueSubscribe(n.Cfg.Subject, n.Cfg.Queue, msgChan)
	if err != nil {
		n.logger.Printf("%s failed to create NATS subscription: %v", workerLogPrefix, err)
		time.Sleep(n.Cfg.ConnectTimeWait)
		nc.Close()
		goto START
	}
	defer close(msgChan)
	defer sub.Unsubscribe()

	for {
		select {
		case <-ctx.Done():
			return
		case m, ok := <-msgChan:
			if !ok {
				n.logger.Printf("%s channel closed, retrying...", workerLogPrefix)
				time.Sleep(n.Cfg.ConnectTimeWait)
				nc.Close()
				goto START
			}
			if n.Cfg.Debug {
				n.logger.Printf("received msg, len=%d, data=%s", len(m.Data), string(m.Data))
			}

			switch n.Cfg.Format {
			case "event":
				evMsgs := make([]*formatters.EventMsg, 1)
				err = json.Unmarshal(m.Data, &evMsgs)
				if err != nil {
					if n.Cfg.Debug {
						n.logger.Printf("%s failed to unmarshal event msg: %v", workerLogPrefix, err)
					}
					continue
				}
				go func() {
					for _, o := range n.outputs {
						for _, ev := range evMsgs {
							o.WriteEvent(ctx, ev)
						}
					}
				}()
			case "proto":
				var protoMsg proto.Message
				err = proto.Unmarshal(m.Data, protoMsg)
				if err != nil {
					if n.Cfg.Debug {
						n.logger.Printf("failed to unmarshal proto msg: %v", err)
					}
					continue
				}
				meta := outputs.Meta{}
				subjectSections := strings.SplitN(m.Subject, ".", 3)
				if len(subjectSections) == 3 {
					meta["source"] = strings.ReplaceAll(subjectSections[1], "-", ".")
					meta["subscription-name"] = subjectSections[2]
				}
				go func() {
					for _, o := range n.outputs {
						o.Write(ctx, protoMsg, meta)
					}
				}()
			}

		}
	}
}

// Close //
func (n *NatsInput) Close() error {
	n.cfn()
	n.wg.Wait()
	return nil
}

// SetLogger //
func (n *NatsInput) SetLogger(logger *log.Logger) {
	if logger != nil && n.logger != nil {
		n.logger.SetOutput(logger.Writer())
		n.logger.SetFlags(logger.Flags())
	}
}

// SetOutputs //
func (n *NatsInput) SetOutputs(outs map[string]outputs.Output) {
	if len(n.Cfg.Outputs) == 0 {
		for _, o := range outs {
			n.outputs = append(n.outputs, o)
		}
		return
	}
	for _, name := range n.Cfg.Outputs {
		if o, ok := outs[name]; ok {
			n.outputs = append(n.outputs, o)
		}
	}
}

// helper functions

func (n *NatsInput) setDefaults() error {
	if n.Cfg.Format == "" {
		n.Cfg.Format = defaultFormat
	}
	if !(strings.ToLower(n.Cfg.Format) == "event" || strings.ToLower(n.Cfg.Format) == "proto") {
		return fmt.Errorf("unsupported input format")
	}
	if n.Cfg.Name == "" {
		n.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	if n.Cfg.Subject == "" {
		n.Cfg.Subject = defaultSubject
	}
	if n.Cfg.Address == "" {
		n.Cfg.Address = defaultAddress
	}
	if n.Cfg.ConnectTimeWait <= 0 {
		n.Cfg.ConnectTimeWait = natsConnectWait
	}
	if n.Cfg.Queue == "" {
		n.Cfg.Queue = n.Cfg.Name
	}
	if n.Cfg.NumWorkers <= 0 {
		n.Cfg.NumWorkers = defaultNumWorkers
	}
	return nil
}

func (n *NatsInput) createNATSConn(c *Config) (*nats.Conn, error) {
	opts := []nats.Option{
		nats.Name(c.Name),
		nats.SetCustomDialer(n),
		nats.ReconnectWait(n.Cfg.ConnectTimeWait),
		nats.ReconnectBufSize(natsReconnectBufferSize),
		nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) {
			n.logger.Printf("NATS error: %v", err)
		}),
		nats.DisconnectHandler(func(c *nats.Conn) {
			n.logger.Println("Disconnected from NATS")
		}),
		nats.ClosedHandler(func(c *nats.Conn) {
			n.logger.Println("NATS connection is closed")
		}),
	}
	if c.Username != "" && c.Password != "" {
		opts = append(opts, nats.UserInfo(c.Username, c.Password))
	}
	nc, err := nats.Connect(c.Address, opts...)
	if err != nil {
		return nil, err
	}
	return nc, nil
}

// Dial //
func (n *NatsInput) Dial(network, address string) (net.Conn, error) {
	ctx, cancel := context.WithCancel(n.ctx)
	defer cancel()

	for {
		n.logger.Printf("attempting to connect to %s", address)
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		select {
		case <-n.ctx.Done():
			return nil, n.ctx.Err()
		default:
			d := &net.Dialer{}
			if conn, err := d.DialContext(ctx, network, address); err == nil {
				n.logger.Printf("successfully connected to NATS server %s", address)
				return conn, nil
			}
			time.Sleep(n.Cfg.ConnectTimeWait)
		}
	}
}
<|endoftext|>


<commit_before>/* For license and copyright information please see LEGAL file in repository */

package assets

import (
	"bytes"
	"compress/gzip"
	"hash/crc32"
	"regexp"
	"strconv"

	"../log"

	"github.com/tdewolff/minify"
	"github.com/tdewolff/minify/css"
	"github.com/tdewolff/minify/html"
	"github.com/tdewolff/minify/js"
	"github.com/tdewolff/minify/json"
	"github.com/tdewolff/minify/svg"
	"github.com/tdewolff/minify/xml"
)

// File store some data about a file!
type File struct {
	FullName     string
	Name         string
	Extension    string
	MimeType     string
	Dep          *Folder
	State        uint8
	Data         []byte
	CompressData []byte
	CompressType string
}

// File||Folder State
const (
	StateUnChanged uint8 = iota
	StateChanged
)

// var md5Hasher = md5.New()

// GetHashOfData return hash of f.Data
func (f *File) GetHashOfData() (hash string) {
	// Just want to differ two same file, So crc32 is more enough!
	// md5Hasher.Write(f.Data)
	// hash = hex.EncodeToString(md5Hasher.Sum(nil))
	// md5Hasher.Reset()
	return strconv.FormatUint(uint64(crc32.ChecksumIEEE(f.Data)), 10)
}

// AddHashToName add f.Data hash to its name and full name.
func (f *File) AddHashToName() {
	f.Name += "-" + f.GetHashOfData()
	f.FullName = f.Name + "." + f.Extension
}

// Copy returns a copy of the file.
func (f *File) Copy() *File {
	var file = File{
		FullName:  f.FullName,
		Name:      f.Name,
		Extension: f.Extension,
		MimeType:  f.MimeType,
		Dep:       f.Dep,
		State:     f.State,
		Data:      f.Data,
	}
	return &file
}

// DeepCopy returns a deep copy of the file.
func (f *File) DeepCopy() *File {
	var file = File{
		FullName:  f.FullName,
		Name:      f.Name,
		Extension: f.Extension,
		MimeType:  f.MimeType,
		Dep:       f.Dep,
		State:     f.State,
	}
	copy(file.Data, f.Data)
	return &file
}

var m = minify.New()
var htmlMinify = regexp.MustCompile(`\r?\n|(<!--.*?-->)|(<!--[\w\W\n\s]+?-->)`)

func init() {
	m.AddFunc("text/css", css.Minify)
	m.AddFunc("text/html", html.Minify)
	m.AddFunc("image/svg+xml", svg.Minify)
	m.AddFuncRegexp(regexp.MustCompile("^(application|text)/(x-)?(java|ecma)script$"), js.Minify)
	m.AddFuncRegexp(regexp.MustCompile("[/+]json$"), json.Minify)
	m.AddFuncRegexp(regexp.MustCompile("[/+]xml$"), xml.Minify)
}

// Minify replace file data with minify of them.
func (f *File) Minify() {
	// TODO::: have problem minify HTML >> https://github.com/tdewolff/minify/issues/318
	if f.Extension == "html" {
		f.Data = htmlMinify.ReplaceAll(f.Data, []byte{})
		return
	}
	var err error
	f.Data, err = m.Bytes(f.MimeType, f.Data)
	if err != nil {
		log.Warn("Minify", f.FullName, "occur this error:", err)
	}
}

// Supported compress types
const (
	CompressTypeGZIP = "gzip"
)

// Compress use to compress file data to mostly to use in serving file by servers.
func (f *File) Compress(compressType string) {
	// Check file type and compress just if it worth.
	switch f.Extension {
	case "png", "jpg", "gif", "jpeg", "mkv", "avi", "mp3", "mp4":
		f.CompressData = f.Data
		return
	}

	f.CompressType = compressType
	switch compressType {
	case CompressTypeGZIP:
		var b bytes.Buffer
		var gz = gzip.NewWriter(&b)
		gz.Write(f.Data)
		gz.Close()
		f.CompressData = b.Bytes()
	}
}

// ReplaceReq is request structure of Replace method.
type ReplaceReq struct {
	Data  string
	Start int
	End   int
}

// Replace replace given data in the file
func (f *File) Replace(data []ReplaceReq) {
	var addedSize int
	for _, d := range data {
		d.Start += addedSize
		d.End += addedSize

		var ln = len(d.Data)
		addedSize = ln - (d.End - d.Start)
		if addedSize > 0 {
			if (cap(f.Data) - len(f.Data)) < addedSize {
				f.Data = append(f.Data, make([]byte, addedSize)...)
			} else {
				// increase f.Data len
				f.Data = f.Data[:len(f.Data)+addedSize]
			}
		}

		copy(f.Data[d.End+addedSize:], f.Data[d.End:])
		copy(f.Data[d.Start:], d.Data)

		if addedSize < 0 {
			// decrease f.Data len
			f.Data = f.Data[:len(f.Data)+addedSize]
		}
	}
}
<commit_msg>Fix Replace() method of File in assets package Previous logic prevent replace work correctly for more than 2 replace data!<commit_after>/* For license and copyright information please see LEGAL file in repository */

package assets

import (
	"bytes"
	"compress/gzip"
	"hash/crc32"
	"regexp"
	"strconv"

	"../log"

	"github.com/tdewolff/minify"
	"github.com/tdewolff/minify/css"
	"github.com/tdewolff/minify/html"
	"github.com/tdewolff/minify/js"
	"github.com/tdewolff/minify/json"
	"github.com/tdewolff/minify/svg"
	"github.com/tdewolff/minify/xml"
)

// File store some data about a file!
type File struct {
	FullName     string
	Name         string
	Extension    string
	MimeType     string
	Dep          *Folder
	State        uint8
	Data         []byte
	CompressData []byte
	CompressType string
}

// File||Folder State
const (
	StateUnChanged uint8 = iota
	StateChanged
)

// var md5Hasher = md5.New()

// GetHashOfData return hash of f.Data
func (f *File) GetHashOfData() (hash string) {
	// Just want to differ two same file, So crc32 is more enough!
	// md5Hasher.Write(f.Data)
	// hash = hex.EncodeToString(md5Hasher.Sum(nil))
	// md5Hasher.Reset()
	return strconv.FormatUint(uint64(crc32.ChecksumIEEE(f.Data)), 10)
}

// AddHashToName add f.Data hash to its name and full name.
func (f *File) AddHashToName() {
	f.Name += "-" + f.GetHashOfData()
	f.FullName = f.Name + "." + f.Extension
}

// Copy returns a copy of the file.
func (f *File) Copy() *File {
	var file = File{
		FullName:  f.FullName,
		Name:      f.Name,
		Extension: f.Extension,
		MimeType:  f.MimeType,
		Dep:       f.Dep,
		State:     f.State,
		Data:      f.Data,
	}
	return &file
}

// DeepCopy returns a deep copy of the file.
func (f *File) DeepCopy() *File {
	var file = File{
		FullName:  f.FullName,
		Name:      f.Name,
		Extension: f.Extension,
		MimeType:  f.MimeType,
		Dep:       f.Dep,
		State:     f.State,
	}
	copy(file.Data, f.Data)
	return &file
}

var m = minify.New()
var htmlMinify = regexp.MustCompile(`\r?\n|(<!--.*?-->)|(<!--[\w\W\n\s]+?-->)`)

func init() {
	m.AddFunc("text/css", css.Minify)
	m.AddFunc("text/html", html.Minify)
	m.AddFunc("image/svg+xml", svg.Minify)
	m.AddFuncRegexp(regexp.MustCompile("^(application|text)/(x-)?(java|ecma)script$"), js.Minify)
	m.AddFuncRegexp(regexp.MustCompile("[/+]json$"), json.Minify)
	m.AddFuncRegexp(regexp.MustCompile("[/+]xml$"), xml.Minify)
}

// Minify replace file data with minify of them.
func (f *File) Minify() {
	// TODO::: have problem minify HTML >> https://github.com/tdewolff/minify/issues/318
	if f.Extension == "html" {
		f.Data = htmlMinify.ReplaceAll(f.Data, []byte{})
		return
	}
	var err error
	f.Data, err = m.Bytes(f.MimeType, f.Data)
	if err != nil {
		log.Warn("Minify", f.FullName, "occur this error:", err)
	}
}

// Supported compress types
const (
	CompressTypeGZIP = "gzip"
)

// Compress use to compress file data to mostly to use in serving file by servers.
func (f *File) Compress(compressType string) {
	// Check file type and compress just if it worth.
	switch f.Extension {
	case "png", "jpg", "gif", "jpeg", "mkv", "avi", "mp3", "mp4":
		f.CompressData = f.Data
		return
	}

	f.CompressType = compressType
	switch compressType {
	case CompressTypeGZIP:
		var b bytes.Buffer
		var gz = gzip.NewWriter(&b)
		gz.Write(f.Data)
		gz.Close()
		f.CompressData = b.Bytes()
	}
}

// ReplaceReq is request structure of Replace method.
type ReplaceReq struct {
	Data  string
	Start int
	End   int
}

// Replace replace given data in the file
func (f *File) Replace(data []ReplaceReq) {
	var totalAddedSize, addedSize int
	for _, d := range data {
		d.Start += totalAddedSize
		d.End += totalAddedSize

		var ln = len(d.Data)
		addedSize = ln - (d.End - d.Start)
		if addedSize > 0 {
			if (cap(f.Data) - len(f.Data)) < addedSize {
				f.Data = append(f.Data, make([]byte, addedSize)...)
			} else {
				// increase f.Data len
				f.Data = f.Data[:len(f.Data)+addedSize]
			}
		}
		totalAddedSize += addedSize

		copy(f.Data[d.End+addedSize:], f.Data[d.End:])
		copy(f.Data[d.Start:], d.Data)

		if addedSize < 0 {
			// decrease f.Data len
			f.Data = f.Data[:len(f.Data)+addedSize]
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
	"io/ioutil"
	. "github.com/onsi/gomega"
	"os"
)

var testDirectory = "/tmp/hoverfly-cli-tests"

func setup() {
	os.Mkdir(testDirectory, 0777)
}

func teardown() {
	os.RemoveAll(testDirectory)
}

func Test_LocalCache_WriteSimulation(t *testing.T) {
	RegisterTestingT(t)
	setup()

	ioutil.WriteFile("/tmp/vendor.name.v1.hfile", []byte("this is a test file"), 0644)

	localCache := LocalCache{Uri: "/tmp"}

	data, err := localCache.ReadSimulation(Hoverfile{Vendor: "vendor", Name: "name", Version: "v1"})

	Expect(err).To(BeNil())
	Expect(data).To(Equal([]byte("this is a test file")))

	teardown()
}<commit_msg>Testing that we throw an appriopriate error if the file does not exist<commit_after>package main

import (
	"testing"
	"io/ioutil"
	. "github.com/onsi/gomega"
	"os"
	"fmt"
)

var testDirectory = "/tmp/hoverfly-cli-tests"

func setup() {
	os.Mkdir(testDirectory, 0777)
}

func teardown() {
	os.RemoveAll(testDirectory)
}

func Test_LocalCache_WriteSimulation(t *testing.T) {
	RegisterTestingT(t)
	setup()

	ioutil.WriteFile("/tmp/vendor.name.v1.hfile", []byte("this is a test file"), 0644)

	localCache := LocalCache{Uri: "/tmp"}

	data, err := localCache.ReadSimulation(Hoverfile{Vendor: "vendor", Name: "name", Version: "v1"})

	Expect(err).To(BeNil())
	Expect(data).To(Equal([]byte("this is a test file")))

	teardown()
}

func Test_LocalCache_WriteSimulation_ErrorsWhenFileIsMissing(t *testing.T) {
	RegisterTestingT(t)
	setup()

	localCache := LocalCache{Uri: "/tmp"}

	data, err := localCache.ReadSimulation(Hoverfile{Vendor: "vendor", Name: "name", Version: "v1"})
	fmt.Println(string(data))
	Expect(err).ToNot(BeNil())
	Expect(err.Error()).To(Equal("Simulation not found"))
	Expect(data).To(BeNil())

	teardown()
}<|endoftext|>


<commit_before>package ngorm

import (
	"errors"
	"fmt"
	"reflect"

	"github.com/ngorm/ngorm/model"
	"github.com/ngorm/ngorm/scope"
	"github.com/ngorm/ngorm/util"
)

// Association provides utility functions for dealing with association queries
type Association struct {
	db     *DB
	column string
	field  *model.Field
}

// Find find out all related associations
func (a *Association) Find(v interface{}) error {
	return a.db.Related(v, a.column)
}

// Append append new associations for many2many, has_many, replace current association for has_one, belongs_to
func (a *Association) Append(values ...interface{}) error {
	if rel := a.field.Relationship; rel.Kind == "has_one" {
		return a.Replace(values...)
	}
	return a.Save(values...)
}

// Replace replace current associations with new one
func (a *Association) Replace(values ...interface{}) error {
	e := a.db.e
	rel := a.field.Relationship
	ndb := a.db.Begin()
	a.field.Set(reflect.Zero(a.field.Field.Type()))
	if err := a.Save(values...); err != nil {
		return err
	}
	switch rel.Kind {
	case "belongs_to":
		if len(values) == 0 {
			m := make(map[string]interface{})
			for _, foreignKey := range rel.ForeignDBNames {
				m[foreignKey] = nil
			}
			return a.db.Begin().Model(e.Scope.Value).UpdateColumn(m)
		}
	default:
		if rel.PolymorphicDBName != "" {
			ndb = ndb.Where(fmt.Sprintf("%v = ?",
				scope.Quote(e, rel.PolymorphicDBName)), rel.PolymorphicValue)
		}
		// Delete Relations except new created
		if len(values) > 0 {
			var fNames, fdDBNames []string
			if rel.Kind == "many_to_many" {
				// if many to many relations, get association fields name from association foreign keys
				se := e.Clone()
				v := reflect.New(a.field.Field.Type()).Interface()
				for idx, dbName := range rel.AssociationForeignFieldNames {
					if field, err := scope.FieldByName(se, v, dbName); err == nil {
						fNames = append(fNames, field.Name)
						fdDBNames = append(fdDBNames, rel.AssociationForeignDBNames[idx])
					}
				}
			} else {
				// If has one/many relations, use primary keys
				fds, err := scope.PrimaryFields(e.Clone(), reflect.New(a.field.Field.Type()).Interface())
				if err != nil {
					return err
				}
				for _, field := range fds {
					fNames = append(fNames, field.Name)
					fdDBNames = append(fdDBNames, field.DBName)
				}
			}

			newPrimaryKeys := util.ColumnAsArray(fNames, a.field.Field.Interface())

			if len(newPrimaryKeys) > 0 {
				sql := fmt.Sprintf("%v NOT IN (%v)",
					scope.ToQueryCondition(ndb.e, fdDBNames),
					scope.ToQueryMarks(newPrimaryKeys))
				ndb = ndb.Where(sql, util.ToQueryValues(newPrimaryKeys)...)
			}
		}
		if rel.Kind == "many_to_many" {
			// if many to many relations, delete related relations from join table
			var sourceForeignFieldNames []string

			for _, dbName := range rel.ForeignFieldNames {
				if field, err := scope.FieldByName(e, e.Scope.Value, dbName); err == nil {
					sourceForeignFieldNames = append(sourceForeignFieldNames, field.Name)
				}
			}

			if sourcePrimaryKeys := util.ColumnAsArray(sourceForeignFieldNames, e.Scope.Value); len(sourcePrimaryKeys) > 0 {
				ndb = ndb.Where(fmt.Sprintf("%v IN (%v)",
					scope.ToQueryCondition(e, rel.ForeignDBNames),
					scope.ToQueryMarks(sourcePrimaryKeys)),
					util.ToQueryValues(sourcePrimaryKeys)...)

				// association.setErr(relationship.JoinTableHandler.Delete(relationship.JoinTableHandler, newDB, relationship))
			}
		} else if rel.Kind == "has_one" || rel.Kind == "has_many" {
			// has_one or has_many relations, set foreign key to be nil (TODO or delete them?)
			var foreignKeyMap = map[string]interface{}{}
			for idx, foreignKey := range rel.ForeignDBNames {
				foreignKeyMap[foreignKey] = nil
				if field, err := scope.FieldByName(e, e.Scope.Value, rel.AssociationForeignFieldNames[idx]); err == nil {
					ndb = ndb.Where(fmt.Sprintf("%v = ?",
						scope.Quote(e, foreignKey)), field.Field.Interface())
				}
			}

			fieldValue := reflect.New(a.field.Field.Type()).Interface()
			return ndb.Model(fieldValue).UpdateColumn(foreignKeyMap)
		}

	}

	return nil
}

// Save save passed values as associations
func (a *Association) Save(values ...interface{}) error {
	rel := a.field.Relationship
	e := a.db.e
	field := a.field

	saveAssociation := func(reflectValue reflect.Value) error {
		// value has to been pointer
		if reflectValue.Kind() != reflect.Ptr {
			reflectPtr := reflect.New(reflectValue.Type())
			reflectPtr.Elem().Set(reflectValue)
			reflectValue = reflectPtr
		}

		// value has to been saved for many2many
		if rel.Kind == "many_to_many" {
			_, err := scope.PrimaryField(e.Clone(), reflectValue.Interface())
			if err == nil {
				return a.db.Begin().Save(reflectValue.Interface())
			}
			return nil

		}

		// Assign Fields
		var fieldType = field.Field.Type()
		var setFieldBackToValue, setSliceFieldBackToValue bool
		if reflectValue.Type().AssignableTo(fieldType) {
			field.Set(reflectValue)
		} else if reflectValue.Type().Elem().AssignableTo(fieldType) {
			// if field's type is struct, then need to set value back to argument after save
			setFieldBackToValue = true
			field.Set(reflectValue.Elem())
		} else if fieldType.Kind() == reflect.Slice {
			if reflectValue.Type().AssignableTo(fieldType.Elem()) {
				field.Set(reflect.Append(field.Field, reflectValue))
			} else if reflectValue.Type().Elem().AssignableTo(fieldType.Elem()) {
				// if field's type is slice of struct, then need to set value back to argument after save
				setSliceFieldBackToValue = true
				field.Set(reflect.Append(field.Field, reflectValue.Elem()))
			}
		}

		if rel.Kind == "many_to_many" {
			// association.setErr(relationship.JoinTableHandler.Add(relationship.JoinTableHandler, scope.NewDB(), scope.Value, reflectValue.Interface()))
		} else {
			err := a.db.Begin().Select(field.Name).Save(e.Scope.Value)
			if err != nil {
				return err
			}
			if setFieldBackToValue {
				reflectValue.Elem().Set(field.Field)
			} else if setSliceFieldBackToValue {
				reflectValue.Elem().Set(field.Field.Index(field.Field.Len() - 1))
			}
		}
		return nil
	}

	for _, value := range values {
		reflectValue := reflect.ValueOf(value)
		indirectReflectValue := reflect.Indirect(reflectValue)
		if indirectReflectValue.Kind() == reflect.Struct {
			saveAssociation(reflectValue)
		} else if indirectReflectValue.Kind() == reflect.Slice {
			for i := 0; i < indirectReflectValue.Len(); i++ {
				err := saveAssociation(indirectReflectValue.Index(i))
				if err != nil {
					return err
				}
			}
		} else {
			return errors.New("invalid value type")
		}
	}
	return nil
}

// Count return the count of current associations
func (a *Association) Count() (int, error) {
	var (
		count      = 0
		rel        = a.field.Relationship
		fieldValue = a.field.Field.Interface()
		query      = a.db.Begin().Model(fieldValue)
	)
	if rel.Kind == "many_to_many" {
		err := scope.JoinWith(rel.JoinTableHandler, query.e, a.db.e.Scope.Value)
		if err != nil {
			return 0, err
		}
	} else if rel.Kind == "has_many" || rel.Kind == "has_one" {
		primaryKeys := util.ColumnAsArray(rel.AssociationForeignFieldNames, a.db.e.Scope.Value)
		query = query.Where(
			fmt.Sprintf("%v IN (%v)",
				scope.ToQueryCondition(a.db.e, rel.ForeignDBNames),
				scope.ToQueryMarks(primaryKeys)),
			util.ToQueryValues(primaryKeys)...,
		)
	} else if rel.Kind == "belongs_to" {
		primaryKeys := util.ColumnAsArray(rel.ForeignFieldNames, a.db.e.Scope.Value)
		query = query.Where(
			fmt.Sprintf("%v IN (%v)",
				scope.ToQueryCondition(a.db.e, rel.AssociationForeignDBNames),
				scope.ToQueryMarks(primaryKeys)),
			util.ToQueryValues(primaryKeys)...,
		)
	}

	if rel.PolymorphicType != "" {
		query = query.Where(
			fmt.Sprintf("%v.%v = ?",
				scope.QuotedTableName(a.db.e, fieldValue),
				scope.Quote(a.db.e, rel.PolymorphicDBName)),
			rel.PolymorphicValue,
		)
	}
	query.e.Scope.Value = fieldValue
	// pretty.Println(query.e.Scope)
	// pretty.Println(query.e.Scope.Value)

	err := query.Count(&count)
	if err != nil {
		return 0, err
	}
	// pretty.Println(query.e.Scope)

	return count, nil
}
<commit_msg>[ngorm] fix field name for ql<commit_after>package ngorm

import (
	"errors"
	"fmt"
	"reflect"

	"github.com/ngorm/ngorm/model"
	"github.com/ngorm/ngorm/scope"
	"github.com/ngorm/ngorm/util"
)

// Association provides utility functions for dealing with association queries
type Association struct {
	db     *DB
	column string
	field  *model.Field
}

// Find find out all related associations
func (a *Association) Find(v interface{}) error {
	return a.db.Related(v, a.column)
}

// Append append new associations for many2many, has_many, replace current association for has_one, belongs_to
func (a *Association) Append(values ...interface{}) error {
	if rel := a.field.Relationship; rel.Kind == "has_one" {
		return a.Replace(values...)
	}
	return a.Save(values...)
}

// Replace replace current associations with new one
func (a *Association) Replace(values ...interface{}) error {
	e := a.db.e
	rel := a.field.Relationship
	ndb := a.db.Begin()
	a.field.Set(reflect.Zero(a.field.Field.Type()))
	if err := a.Save(values...); err != nil {
		return err
	}
	switch rel.Kind {
	case "belongs_to":
		if len(values) == 0 {
			m := make(map[string]interface{})
			for _, foreignKey := range rel.ForeignDBNames {
				m[foreignKey] = nil
			}
			return a.db.Begin().Model(e.Scope.Value).UpdateColumn(m)
		}
	default:
		if rel.PolymorphicDBName != "" {
			ndb = ndb.Where(fmt.Sprintf("%v = ?",
				scope.Quote(e, rel.PolymorphicDBName)), rel.PolymorphicValue)
		}
		// Delete Relations except new created
		if len(values) > 0 {
			var fNames, fdDBNames []string
			if rel.Kind == "many_to_many" {
				// if many to many relations, get association fields name from association foreign keys
				se := e.Clone()
				v := reflect.New(a.field.Field.Type()).Interface()
				for idx, dbName := range rel.AssociationForeignFieldNames {
					if field, err := scope.FieldByName(se, v, dbName); err == nil {
						fNames = append(fNames, field.Name)
						fdDBNames = append(fdDBNames, rel.AssociationForeignDBNames[idx])
					}
				}
			} else {
				// If has one/many relations, use primary keys
				fds, err := scope.PrimaryFields(e.Clone(), reflect.New(a.field.Field.Type()).Interface())
				if err != nil {
					return err
				}
				for _, field := range fds {
					fNames = append(fNames, field.Name)
					fdDBNames = append(fdDBNames, field.DBName)
				}
			}

			newPrimaryKeys := util.ColumnAsArray(fNames, a.field.Field.Interface())

			if len(newPrimaryKeys) > 0 {
				sql := fmt.Sprintf("%v NOT IN (%v)",
					scope.ToQueryCondition(ndb.e, fdDBNames),
					scope.ToQueryMarks(newPrimaryKeys))
				ndb = ndb.Where(sql, util.ToQueryValues(newPrimaryKeys)...)
			}
		}
		if rel.Kind == "many_to_many" {
			// if many to many relations, delete related relations from join table
			var sourceForeignFieldNames []string

			for _, dbName := range rel.ForeignFieldNames {
				if field, err := scope.FieldByName(e, e.Scope.Value, dbName); err == nil {
					sourceForeignFieldNames = append(sourceForeignFieldNames, field.Name)
				}
			}

			if sourcePrimaryKeys := util.ColumnAsArray(sourceForeignFieldNames, e.Scope.Value); len(sourcePrimaryKeys) > 0 {
				ndb = ndb.Where(fmt.Sprintf("%v IN (%v)",
					scope.ToQueryCondition(e, rel.ForeignDBNames),
					scope.ToQueryMarks(sourcePrimaryKeys)),
					util.ToQueryValues(sourcePrimaryKeys)...)

				// association.setErr(relationship.JoinTableHandler.Delete(relationship.JoinTableHandler, newDB, relationship))
			}
		} else if rel.Kind == "has_one" || rel.Kind == "has_many" {
			// has_one or has_many relations, set foreign key to be nil (TODO or delete them?)
			var foreignKeyMap = map[string]interface{}{}
			for idx, foreignKey := range rel.ForeignDBNames {
				foreignKeyMap[foreignKey] = nil
				if field, err := scope.FieldByName(e, e.Scope.Value, rel.AssociationForeignFieldNames[idx]); err == nil {
					ndb = ndb.Where(fmt.Sprintf("%v = ?",
						scope.Quote(e, foreignKey)), field.Field.Interface())
				}
			}

			fieldValue := reflect.New(a.field.Field.Type()).Interface()
			return ndb.Model(fieldValue).UpdateColumn(foreignKeyMap)
		}

	}

	return nil
}

// Save save passed values as associations
func (a *Association) Save(values ...interface{}) error {
	rel := a.field.Relationship
	e := a.db.e
	field := a.field

	saveAssociation := func(reflectValue reflect.Value) error {
		// value has to been pointer
		if reflectValue.Kind() != reflect.Ptr {
			reflectPtr := reflect.New(reflectValue.Type())
			reflectPtr.Elem().Set(reflectValue)
			reflectValue = reflectPtr
		}

		// value has to been saved for many2many
		if rel.Kind == "many_to_many" {
			_, err := scope.PrimaryField(e.Clone(), reflectValue.Interface())
			if err == nil {
				return a.db.Begin().Save(reflectValue.Interface())
			}
			return nil

		}

		// Assign Fields
		var fieldType = field.Field.Type()
		var setFieldBackToValue, setSliceFieldBackToValue bool
		if reflectValue.Type().AssignableTo(fieldType) {
			field.Set(reflectValue)
		} else if reflectValue.Type().Elem().AssignableTo(fieldType) {
			// if field's type is struct, then need to set value back to argument after save
			setFieldBackToValue = true
			field.Set(reflectValue.Elem())
		} else if fieldType.Kind() == reflect.Slice {
			if reflectValue.Type().AssignableTo(fieldType.Elem()) {
				field.Set(reflect.Append(field.Field, reflectValue))
			} else if reflectValue.Type().Elem().AssignableTo(fieldType.Elem()) {
				// if field's type is slice of struct, then need to set value back to argument after save
				setSliceFieldBackToValue = true
				field.Set(reflect.Append(field.Field, reflectValue.Elem()))
			}
		}

		if rel.Kind == "many_to_many" {
			// association.setErr(relationship.JoinTableHandler.Add(relationship.JoinTableHandler, scope.NewDB(), scope.Value, reflectValue.Interface()))
		} else {
			err := a.db.Begin().Select(field.Name).Save(e.Scope.Value)
			if err != nil {
				return err
			}
			if setFieldBackToValue {
				reflectValue.Elem().Set(field.Field)
			} else if setSliceFieldBackToValue {
				reflectValue.Elem().Set(field.Field.Index(field.Field.Len() - 1))
			}
		}
		return nil
	}

	for _, value := range values {
		reflectValue := reflect.ValueOf(value)
		indirectReflectValue := reflect.Indirect(reflectValue)
		if indirectReflectValue.Kind() == reflect.Struct {
			saveAssociation(reflectValue)
		} else if indirectReflectValue.Kind() == reflect.Slice {
			for i := 0; i < indirectReflectValue.Len(); i++ {
				err := saveAssociation(indirectReflectValue.Index(i))
				if err != nil {
					return err
				}
			}
		} else {
			return errors.New("invalid value type")
		}
	}
	return nil
}

// Count return the count of current associations
func (a *Association) Count() (int, error) {
	var (
		count      = 0
		rel        = a.field.Relationship
		fieldValue = a.field.Field.Interface()
		query      = a.db.Begin().Model(fieldValue)
	)
	if rel.Kind == "many_to_many" {
		err := scope.JoinWith(rel.JoinTableHandler, query.e, a.db.e.Scope.Value)
		if err != nil {
			return 0, err
		}
	} else if rel.Kind == "has_many" || rel.Kind == "has_one" {
		primaryKeys := util.ColumnAsArray(rel.AssociationForeignFieldNames, a.db.e.Scope.Value)
		query = query.Where(
			fmt.Sprintf("%v IN (%v)",
				scope.ToQueryCondition(a.db.e, rel.ForeignDBNames),
				scope.ToQueryMarks(primaryKeys)),
			util.ToQueryValues(primaryKeys)...,
		)
	} else if rel.Kind == "belongs_to" {
		primaryKeys := util.ColumnAsArray(rel.ForeignFieldNames, a.db.e.Scope.Value)
		query = query.Where(
			fmt.Sprintf("%v IN (%v)",
				scope.ToQueryCondition(a.db.e, rel.AssociationForeignDBNames),
				scope.ToQueryMarks(primaryKeys)),
			util.ToQueryValues(primaryKeys)...,
		)
	}

	if rel.PolymorphicType != "" {
		query = query.Where(
			fmt.Sprintf("%v%v = ?",
				a.db.e.Dialect.QueryFieldName(
					scope.QuotedTableName(a.db.e, fieldValue)),
				scope.Quote(a.db.e, rel.PolymorphicDBName)),
			rel.PolymorphicValue,
		)
	}
	err := query.Count(&count)
	if err != nil {
		return 0, err
	}
	return count, nil
}
<|endoftext|>


<commit_before>package sandbox

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/cri-o/cri-o/internal/config/nsmgr"
	"github.com/cri-o/cri-o/internal/hostport"
	"github.com/cri-o/cri-o/internal/oci"
	"github.com/cri-o/cri-o/server/cri/types"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/sys/unix"
	"k8s.io/apimachinery/pkg/fields"
)

// DevShmPath is the default system wide shared memory path
const DevShmPath = "/dev/shm"

var (
	sbStoppedFilename        = "stopped"
	sbNetworkStoppedFilename = "network-stopped"
)

// Sandbox contains data surrounding kubernetes sandboxes on the server
type Sandbox struct {
	portMappings []*hostport.PortMapping
	createdAt    time.Time
	id           string
	namespace    string
	// OCI pod name (eg "<namespace>-<name>-<attempt>")
	name string
	// Kubernetes pod name (eg, "<name>")
	kubeName       string
	logDir         string
	containers     oci.ContainerStorer
	processLabel   string
	mountLabel     string
	netns          nsmgr.Namespace
	ipcns          nsmgr.Namespace
	utsns          nsmgr.Namespace
	userns         nsmgr.Namespace
	shmPath        string
	cgroupParent   string
	runtimeHandler string
	resolvPath     string
	hostnamePath   string
	hostname       string
	// ipv4 or ipv6 cache
	ips                []string
	seccompProfilePath string
	labels             fields.Set
	annotations        map[string]string
	infraContainer     *oci.Container
	metadata           *Metadata
	nsOpts             *types.NamespaceOption
	stopMutex          sync.RWMutex
	created            bool
	stopped            bool
	networkStopped     bool
	privileged         bool
	hostNetwork        bool
	usernsMode         string
}

type Metadata struct {
	// Pod name of the sandbox.
	Name string `json:"name,omitempty"`

	// Pod UID of the sandbox.
	UID string `json:"uid,omitempty"`

	// Pod namespace of the sandbox.
	Namespace string `json:"namespace,omitempty"`

	// Attempt number of creating the sandbox.
	Attempt uint32 `json:"attempt,omitempty"`
}

// DefaultShmSize is the default shm size
const DefaultShmSize = 64 * 1024 * 1024

// ErrIDEmpty is the error returned when the id of the sandbox is empty
var ErrIDEmpty = errors.New("PodSandboxId should not be empty")

// New creates and populates a new pod sandbox
// New sandboxes have no containers, no infra container, and no network namespaces associated with them
// An infra container must be attached before the sandbox is added to the state
func New(id, namespace, name, kubeName, logDir string, labels, annotations map[string]string, processLabel, mountLabel string, metadata *Metadata, shmPath, cgroupParent string, privileged bool, runtimeHandler, resolvPath, hostname string, portMappings []*hostport.PortMapping, hostNetwork bool, createdAt time.Time, usernsMode string) (*Sandbox, error) {
	sb := new(Sandbox)
	sb.id = id
	sb.namespace = namespace
	sb.name = name
	sb.kubeName = kubeName
	sb.logDir = logDir
	sb.labels = labels
	sb.annotations = annotations
	sb.containers = oci.NewMemoryStore()
	sb.processLabel = processLabel
	sb.mountLabel = mountLabel
	sb.metadata = metadata
	sb.shmPath = shmPath
	sb.cgroupParent = cgroupParent
	sb.privileged = privileged
	sb.runtimeHandler = runtimeHandler
	sb.resolvPath = resolvPath
	sb.hostname = hostname
	sb.portMappings = portMappings
	sb.createdAt = createdAt
	sb.hostNetwork = hostNetwork
	sb.usernsMode = usernsMode

	return sb, nil
}

func (s *Sandbox) CreatedAt() time.Time {
	return s.createdAt
}

// SetSeccompProfilePath sets the seccomp profile path
func (s *Sandbox) SetSeccompProfilePath(pp string) {
	s.seccompProfilePath = pp
}

// SeccompProfilePath returns the seccomp profile path
func (s *Sandbox) SeccompProfilePath() string {
	return s.seccompProfilePath
}

// AddIPs stores the ip in the sandbox
func (s *Sandbox) AddIPs(ips []string) {
	s.ips = ips
}

// SetNamespaceOptions sets whether the pod is running using host network
func (s *Sandbox) SetNamespaceOptions(nsOpts *types.NamespaceOption) {
	s.nsOpts = nsOpts
}

// NamespaceOptions returns the namespace options for the sandbox
func (s *Sandbox) NamespaceOptions() *types.NamespaceOption {
	return s.nsOpts
}

// StopMutex returns the mutex to use when stopping the sandbox
func (s *Sandbox) StopMutex() *sync.RWMutex {
	return &s.stopMutex
}

// IPs returns the ip of the sandbox
func (s *Sandbox) IPs() []string {
	return s.ips
}

// ID returns the id of the sandbox
func (s *Sandbox) ID() string {
	return s.id
}

// UsernsMode returns the mode for setting the user namespace, if any.
func (s *Sandbox) UsernsMode() string {
	return s.usernsMode
}

// Namespace returns the namespace for the sandbox
func (s *Sandbox) Namespace() string {
	return s.namespace
}

// Name returns the name of the sandbox
func (s *Sandbox) Name() string {
	return s.name
}

// KubeName returns the kubernetes name for the sandbox
func (s *Sandbox) KubeName() string {
	return s.kubeName
}

// LogDir returns the location of the logging directory for the sandbox
func (s *Sandbox) LogDir() string {
	return s.logDir
}

// Labels returns the labels associated with the sandbox
func (s *Sandbox) Labels() fields.Set {
	return s.labels
}

// Annotations returns a list of annotations for the sandbox
func (s *Sandbox) Annotations() map[string]string {
	return s.annotations
}

// Containers returns the ContainerStorer that contains information on all
// of the containers in the sandbox
func (s *Sandbox) Containers() oci.ContainerStorer {
	return s.containers
}

// ProcessLabel returns the process label for the sandbox
func (s *Sandbox) ProcessLabel() string {
	return s.processLabel
}

// MountLabel returns the mount label for the sandbox
func (s *Sandbox) MountLabel() string {
	return s.mountLabel
}

// Metadata returns a set of metadata about the sandbox
func (s *Sandbox) Metadata() *Metadata {
	return s.metadata
}

// ShmPath returns the shm path of the sandbox
func (s *Sandbox) ShmPath() string {
	return s.shmPath
}

// CgroupParent returns the cgroup parent of the sandbox
func (s *Sandbox) CgroupParent() string {
	return s.cgroupParent
}

// Privileged returns whether or not the containers in the sandbox are
// privileged containers
func (s *Sandbox) Privileged() bool {
	return s.privileged
}

// RuntimeHandler returns the name of the runtime handler that should be
// picked from the list of runtimes. The name must match the key from the
// map of runtimes.
func (s *Sandbox) RuntimeHandler() string {
	return s.runtimeHandler
}

// HostNetwork returns whether the sandbox runs in the host network namespace
func (s *Sandbox) HostNetwork() bool {
	return s.hostNetwork
}

// ResolvPath returns the resolv path for the sandbox
func (s *Sandbox) ResolvPath() string {
	return s.resolvPath
}

// AddHostnamePath adds the hostname path to the sandbox
func (s *Sandbox) AddHostnamePath(hostname string) {
	s.hostnamePath = hostname
}

// HostnamePath retrieves the hostname path from a sandbox
func (s *Sandbox) HostnamePath() string {
	return s.hostnamePath
}

// Hostname returns the hostname of the sandbox
func (s *Sandbox) Hostname() string {
	return s.hostname
}

// PortMappings returns a list of port mappings between the host and the sandbox
func (s *Sandbox) PortMappings() []*hostport.PortMapping {
	return s.portMappings
}

// AddContainer adds a container to the sandbox
func (s *Sandbox) AddContainer(c *oci.Container) {
	s.containers.Add(c.Name(), c)
}

// GetContainer retrieves a container from the sandbox
func (s *Sandbox) GetContainer(name string) *oci.Container {
	return s.containers.Get(name)
}

// RemoveContainer deletes a container from the sandbox
func (s *Sandbox) RemoveContainer(c *oci.Container) {
	s.containers.Delete(c.Name())
}

// SetInfraContainer sets the infrastructure container of a sandbox
// Attempts to set the infrastructure container after one is already present will throw an error
func (s *Sandbox) SetInfraContainer(infraCtr *oci.Container) error {
	if s.infraContainer != nil {
		return fmt.Errorf("sandbox already has an infra container")
	} else if infraCtr == nil {
		return fmt.Errorf("must provide non-nil infra container")
	}

	s.infraContainer = infraCtr

	return nil
}

// InfraContainer returns the infrastructure container for the sandbox
func (s *Sandbox) InfraContainer() *oci.Container {
	return s.infraContainer
}

// RemoveInfraContainer removes the infrastructure container of a sandbox
func (s *Sandbox) RemoveInfraContainer() {
	s.infraContainer = nil
}

// SetStopped sets the sandbox state to stopped.
// This should be set after a stop operation succeeds
// so that subsequent stops can return fast.
// if createFile is true, it also creates a "stopped" file in the infra container's persistent dir
// this is used to track the sandbox is stopped over reboots
func (s *Sandbox) SetStopped(createFile bool) {
	if s.stopped {
		return
	}
	s.stopped = true
	if createFile {
		if err := s.createFileInInfraDir(sbStoppedFilename); err != nil {
			logrus.Errorf("failed to create stopped file in container state. Restore may fail: %v", err)
		}
	}
}

// Stopped returns whether the sandbox state has been
// set to stopped.
func (s *Sandbox) Stopped() bool {
	return s.stopped
}

// SetCreated sets the created status of sandbox to true
func (s *Sandbox) SetCreated() {
	s.created = true
}

// NetworkStopped returns whether the network has been stopped
func (s *Sandbox) NetworkStopped() bool {
	return s.networkStopped
}

// SetNetworkStopped sets the sandbox network state as stopped
// This should be set after a network stop operation succeeds,
// so we don't double stop the network
// if createFile is true, it creates a "network-stopped" file
// in the infra container's persistent dir
// this is used to track the network is stopped over reboots
// returns an error if an error occurred when creating the network-stopped file
func (s *Sandbox) SetNetworkStopped(createFile bool) error {
	if s.networkStopped {
		return nil
	}
	s.networkStopped = true
	if createFile {
		if err := s.createFileInInfraDir(sbNetworkStoppedFilename); err != nil {
			return fmt.Errorf("failed to create state file in container directory. Restores may fail: %v", err)
		}
	}
	return nil
}

func (s *Sandbox) createFileInInfraDir(filename string) error {
	// If the sandbox is not yet created,
	// this function is being called when
	// cleaning up a failed sandbox creation.
	// We don't need to create the file, as there will be no
	// sandbox to restore
	if !s.created {
		return nil
	}
	infra := s.InfraContainer()
	f, err := os.Create(filepath.Join(infra.Dir(), filename))
	if err == nil {
		f.Close()
	}
	return err
}

func (s *Sandbox) RestoreStopped() {
	if s.fileExistsInInfraDir(sbStoppedFilename) {
		s.stopped = true
	}
	if s.fileExistsInInfraDir(sbNetworkStoppedFilename) {
		s.networkStopped = true
	}
}

func (s *Sandbox) fileExistsInInfraDir(filename string) bool {
	infra := s.InfraContainer()
	infraFilePath := filepath.Join(infra.Dir(), filename)
	if _, err := os.Stat(infraFilePath); err != nil {
		if !os.IsNotExist(err) {
			logrus.Warnf("error checking if %s exists: %v", infraFilePath, err)
		}
		return false
	}
	return true
}

// Created returns the created status of sandbox
func (s *Sandbox) Created() bool {
	return s.created
}

// Ready returns whether the sandbox should be marked as ready to the kubelet
// if there is no infra container, it is always considered ready.
// `takeLock` should be set if we need to take the lock to get the infra container's state.
// If there is no infra container, it is never considered ready.
// If the infra container is spoofed, the pod is considered ready when it has been created, but not stopped.
func (s *Sandbox) Ready(takeLock bool) bool {
	podInfraContainer := s.InfraContainer()
	if podInfraContainer == nil {
		return false
	}
	if podInfraContainer.Spoofed() {
		return s.created && !s.stopped
	}
	// Assume the sandbox is ready, unless it has an infra container that
	// isn't running
	var cState *oci.ContainerState
	if takeLock {
		cState = podInfraContainer.State()
	} else {
		cState = podInfraContainer.StateNoLock()
	}

	return cState.Status == oci.ContainerStateRunning
}

// UnmountShm removes the shared memory mount for the sandbox and returns an
// error if any failure occurs.
func (s *Sandbox) UnmountShm() error {
	fp := s.ShmPath()
	if fp == DevShmPath {
		return nil
	}

	// try to unmount, ignoring "not mounted" (EINVAL) error and
	// "already unmounted" (ENOENT) error
	if err := unix.Unmount(fp, unix.MNT_DETACH); err != nil && err != unix.EINVAL && err != unix.ENOENT {
		return errors.Wrapf(err, "unable to unmount %s", fp)
	}

	return nil
}

// NeedsInfra is a function that returns whether the sandbox will need an infra container.
// If the server manages the namespace lifecycles, and the Pid option on the sandbox
// is node or container level, the infra container is not needed
func (s *Sandbox) NeedsInfra(serverDropsInfra bool) bool {
	return !serverDropsInfra || s.nsOpts.Pid == types.NamespaceModePOD
}
<commit_msg>sandbox: fix race with cleanup<commit_after>package sandbox

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/cri-o/cri-o/internal/config/nsmgr"
	"github.com/cri-o/cri-o/internal/hostport"
	"github.com/cri-o/cri-o/internal/oci"
	"github.com/cri-o/cri-o/server/cri/types"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/sys/unix"
	"k8s.io/apimachinery/pkg/fields"
)

// DevShmPath is the default system wide shared memory path
const DevShmPath = "/dev/shm"

var (
	sbStoppedFilename        = "stopped"
	sbNetworkStoppedFilename = "network-stopped"
)

// Sandbox contains data surrounding kubernetes sandboxes on the server
type Sandbox struct {
	portMappings []*hostport.PortMapping
	createdAt    time.Time
	id           string
	namespace    string
	// OCI pod name (eg "<namespace>-<name>-<attempt>")
	name string
	// Kubernetes pod name (eg, "<name>")
	kubeName       string
	logDir         string
	containers     oci.ContainerStorer
	processLabel   string
	mountLabel     string
	netns          nsmgr.Namespace
	ipcns          nsmgr.Namespace
	utsns          nsmgr.Namespace
	userns         nsmgr.Namespace
	shmPath        string
	cgroupParent   string
	runtimeHandler string
	resolvPath     string
	hostnamePath   string
	hostname       string
	// ipv4 or ipv6 cache
	ips                []string
	seccompProfilePath string
	labels             fields.Set
	annotations        map[string]string
	infraContainer     *oci.Container
	metadata           *Metadata
	nsOpts             *types.NamespaceOption
	stopMutex          sync.RWMutex
	created            bool
	stopped            bool
	networkStopped     bool
	privileged         bool
	hostNetwork        bool
	usernsMode         string
}

type Metadata struct {
	// Pod name of the sandbox.
	Name string `json:"name,omitempty"`

	// Pod UID of the sandbox.
	UID string `json:"uid,omitempty"`

	// Pod namespace of the sandbox.
	Namespace string `json:"namespace,omitempty"`

	// Attempt number of creating the sandbox.
	Attempt uint32 `json:"attempt,omitempty"`
}

// DefaultShmSize is the default shm size
const DefaultShmSize = 64 * 1024 * 1024

// ErrIDEmpty is the error returned when the id of the sandbox is empty
var ErrIDEmpty = errors.New("PodSandboxId should not be empty")

// New creates and populates a new pod sandbox
// New sandboxes have no containers, no infra container, and no network namespaces associated with them
// An infra container must be attached before the sandbox is added to the state
func New(id, namespace, name, kubeName, logDir string, labels, annotations map[string]string, processLabel, mountLabel string, metadata *Metadata, shmPath, cgroupParent string, privileged bool, runtimeHandler, resolvPath, hostname string, portMappings []*hostport.PortMapping, hostNetwork bool, createdAt time.Time, usernsMode string) (*Sandbox, error) {
	sb := new(Sandbox)
	sb.id = id
	sb.namespace = namespace
	sb.name = name
	sb.kubeName = kubeName
	sb.logDir = logDir
	sb.labels = labels
	sb.annotations = annotations
	sb.containers = oci.NewMemoryStore()
	sb.processLabel = processLabel
	sb.mountLabel = mountLabel
	sb.metadata = metadata
	sb.shmPath = shmPath
	sb.cgroupParent = cgroupParent
	sb.privileged = privileged
	sb.runtimeHandler = runtimeHandler
	sb.resolvPath = resolvPath
	sb.hostname = hostname
	sb.portMappings = portMappings
	sb.createdAt = createdAt
	sb.hostNetwork = hostNetwork
	sb.usernsMode = usernsMode

	return sb, nil
}

func (s *Sandbox) CreatedAt() time.Time {
	return s.createdAt
}

// SetSeccompProfilePath sets the seccomp profile path
func (s *Sandbox) SetSeccompProfilePath(pp string) {
	s.seccompProfilePath = pp
}

// SeccompProfilePath returns the seccomp profile path
func (s *Sandbox) SeccompProfilePath() string {
	return s.seccompProfilePath
}

// AddIPs stores the ip in the sandbox
func (s *Sandbox) AddIPs(ips []string) {
	s.ips = ips
}

// SetNamespaceOptions sets whether the pod is running using host network
func (s *Sandbox) SetNamespaceOptions(nsOpts *types.NamespaceOption) {
	s.nsOpts = nsOpts
}

// NamespaceOptions returns the namespace options for the sandbox
func (s *Sandbox) NamespaceOptions() *types.NamespaceOption {
	return s.nsOpts
}

// StopMutex returns the mutex to use when stopping the sandbox
func (s *Sandbox) StopMutex() *sync.RWMutex {
	return &s.stopMutex
}

// IPs returns the ip of the sandbox
func (s *Sandbox) IPs() []string {
	return s.ips
}

// ID returns the id of the sandbox
func (s *Sandbox) ID() string {
	return s.id
}

// UsernsMode returns the mode for setting the user namespace, if any.
func (s *Sandbox) UsernsMode() string {
	return s.usernsMode
}

// Namespace returns the namespace for the sandbox
func (s *Sandbox) Namespace() string {
	return s.namespace
}

// Name returns the name of the sandbox
func (s *Sandbox) Name() string {
	return s.name
}

// KubeName returns the kubernetes name for the sandbox
func (s *Sandbox) KubeName() string {
	return s.kubeName
}

// LogDir returns the location of the logging directory for the sandbox
func (s *Sandbox) LogDir() string {
	return s.logDir
}

// Labels returns the labels associated with the sandbox
func (s *Sandbox) Labels() fields.Set {
	return s.labels
}

// Annotations returns a list of annotations for the sandbox
func (s *Sandbox) Annotations() map[string]string {
	return s.annotations
}

// Containers returns the ContainerStorer that contains information on all
// of the containers in the sandbox
func (s *Sandbox) Containers() oci.ContainerStorer {
	return s.containers
}

// ProcessLabel returns the process label for the sandbox
func (s *Sandbox) ProcessLabel() string {
	return s.processLabel
}

// MountLabel returns the mount label for the sandbox
func (s *Sandbox) MountLabel() string {
	return s.mountLabel
}

// Metadata returns a set of metadata about the sandbox
func (s *Sandbox) Metadata() *Metadata {
	return s.metadata
}

// ShmPath returns the shm path of the sandbox
func (s *Sandbox) ShmPath() string {
	return s.shmPath
}

// CgroupParent returns the cgroup parent of the sandbox
func (s *Sandbox) CgroupParent() string {
	return s.cgroupParent
}

// Privileged returns whether or not the containers in the sandbox are
// privileged containers
func (s *Sandbox) Privileged() bool {
	return s.privileged
}

// RuntimeHandler returns the name of the runtime handler that should be
// picked from the list of runtimes. The name must match the key from the
// map of runtimes.
func (s *Sandbox) RuntimeHandler() string {
	return s.runtimeHandler
}

// HostNetwork returns whether the sandbox runs in the host network namespace
func (s *Sandbox) HostNetwork() bool {
	return s.hostNetwork
}

// ResolvPath returns the resolv path for the sandbox
func (s *Sandbox) ResolvPath() string {
	return s.resolvPath
}

// AddHostnamePath adds the hostname path to the sandbox
func (s *Sandbox) AddHostnamePath(hostname string) {
	s.hostnamePath = hostname
}

// HostnamePath retrieves the hostname path from a sandbox
func (s *Sandbox) HostnamePath() string {
	return s.hostnamePath
}

// Hostname returns the hostname of the sandbox
func (s *Sandbox) Hostname() string {
	return s.hostname
}

// PortMappings returns a list of port mappings between the host and the sandbox
func (s *Sandbox) PortMappings() []*hostport.PortMapping {
	return s.portMappings
}

// AddContainer adds a container to the sandbox
func (s *Sandbox) AddContainer(c *oci.Container) {
	s.containers.Add(c.Name(), c)
}

// GetContainer retrieves a container from the sandbox
func (s *Sandbox) GetContainer(name string) *oci.Container {
	return s.containers.Get(name)
}

// RemoveContainer deletes a container from the sandbox
func (s *Sandbox) RemoveContainer(c *oci.Container) {
	s.containers.Delete(c.Name())
}

// SetInfraContainer sets the infrastructure container of a sandbox
// Attempts to set the infrastructure container after one is already present will throw an error
func (s *Sandbox) SetInfraContainer(infraCtr *oci.Container) error {
	if s.infraContainer != nil {
		return fmt.Errorf("sandbox already has an infra container")
	} else if infraCtr == nil {
		return fmt.Errorf("must provide non-nil infra container")
	}

	s.infraContainer = infraCtr

	return nil
}

// InfraContainer returns the infrastructure container for the sandbox
func (s *Sandbox) InfraContainer() *oci.Container {
	return s.infraContainer
}

// RemoveInfraContainer removes the infrastructure container of a sandbox
func (s *Sandbox) RemoveInfraContainer() {
	s.infraContainer = nil
}

// SetStopped sets the sandbox state to stopped.
// This should be set after a stop operation succeeds
// so that subsequent stops can return fast.
// if createFile is true, it also creates a "stopped" file in the infra container's persistent dir
// this is used to track the sandbox is stopped over reboots
func (s *Sandbox) SetStopped(createFile bool) {
	if s.stopped {
		return
	}
	s.stopped = true
	if createFile {
		if err := s.createFileInInfraDir(sbStoppedFilename); err != nil {
			logrus.Errorf("failed to create stopped file in container state. Restore may fail: %v", err)
		}
	}
}

// Stopped returns whether the sandbox state has been
// set to stopped.
func (s *Sandbox) Stopped() bool {
	return s.stopped
}

// SetCreated sets the created status of sandbox to true
func (s *Sandbox) SetCreated() {
	s.created = true
}

// NetworkStopped returns whether the network has been stopped
func (s *Sandbox) NetworkStopped() bool {
	return s.networkStopped
}

// SetNetworkStopped sets the sandbox network state as stopped
// This should be set after a network stop operation succeeds,
// so we don't double stop the network
// if createFile is true, it creates a "network-stopped" file
// in the infra container's persistent dir
// this is used to track the network is stopped over reboots
// returns an error if an error occurred when creating the network-stopped file
func (s *Sandbox) SetNetworkStopped(createFile bool) error {
	if s.networkStopped {
		return nil
	}
	s.networkStopped = true
	if createFile {
		if err := s.createFileInInfraDir(sbNetworkStoppedFilename); err != nil {
			return fmt.Errorf("failed to create state file in container directory. Restores may fail: %v", err)
		}
	}
	return nil
}

func (s *Sandbox) createFileInInfraDir(filename string) error {
	// If the sandbox is not yet created,
	// this function is being called when
	// cleaning up a failed sandbox creation.
	// We don't need to create the file, as there will be no
	// sandbox to restore
	if !s.created {
		return nil
	}
	infra := s.InfraContainer()
	// If the infra directory has been cleaned up already, we should not fail to
	// create this file.
	if _, err := os.Stat(infra.Dir()); os.IsNotExist(err) {
		return nil
	}
	f, err := os.Create(filepath.Join(infra.Dir(), filename))
	if err == nil {
		f.Close()
	}
	return err
}

func (s *Sandbox) RestoreStopped() {
	if s.fileExistsInInfraDir(sbStoppedFilename) {
		s.stopped = true
	}
	if s.fileExistsInInfraDir(sbNetworkStoppedFilename) {
		s.networkStopped = true
	}
}

func (s *Sandbox) fileExistsInInfraDir(filename string) bool {
	infra := s.InfraContainer()
	infraFilePath := filepath.Join(infra.Dir(), filename)
	if _, err := os.Stat(infraFilePath); err != nil {
		if !os.IsNotExist(err) {
			logrus.Warnf("error checking if %s exists: %v", infraFilePath, err)
		}
		return false
	}
	return true
}

// Created returns the created status of sandbox
func (s *Sandbox) Created() bool {
	return s.created
}

// Ready returns whether the sandbox should be marked as ready to the kubelet
// if there is no infra container, it is always considered ready.
// `takeLock` should be set if we need to take the lock to get the infra container's state.
// If there is no infra container, it is never considered ready.
// If the infra container is spoofed, the pod is considered ready when it has been created, but not stopped.
func (s *Sandbox) Ready(takeLock bool) bool {
	podInfraContainer := s.InfraContainer()
	if podInfraContainer == nil {
		return false
	}
	if podInfraContainer.Spoofed() {
		return s.created && !s.stopped
	}
	// Assume the sandbox is ready, unless it has an infra container that
	// isn't running
	var cState *oci.ContainerState
	if takeLock {
		cState = podInfraContainer.State()
	} else {
		cState = podInfraContainer.StateNoLock()
	}

	return cState.Status == oci.ContainerStateRunning
}

// UnmountShm removes the shared memory mount for the sandbox and returns an
// error if any failure occurs.
func (s *Sandbox) UnmountShm() error {
	fp := s.ShmPath()
	if fp == DevShmPath {
		return nil
	}

	// try to unmount, ignoring "not mounted" (EINVAL) error and
	// "already unmounted" (ENOENT) error
	if err := unix.Unmount(fp, unix.MNT_DETACH); err != nil && err != unix.EINVAL && err != unix.ENOENT {
		return errors.Wrapf(err, "unable to unmount %s", fp)
	}

	return nil
}

// NeedsInfra is a function that returns whether the sandbox will need an infra container.
// If the server manages the namespace lifecycles, and the Pid option on the sandbox
// is node or container level, the infra container is not needed
func (s *Sandbox) NeedsInfra(serverDropsInfra bool) bool {
	return !serverDropsInfra || s.nsOpts.Pid == types.NamespaceModePOD
}
<|endoftext|>