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

// generator scours the github.com/Azure/azure-rest-api-specs repository in search for
// REST functionality that has Go configuration, then generates packages accordingly.
// This tool was developed with the intention that it be an internal tool for the
// Azure-SDK-for-Go team, but the usage of this package for public scenarios is in no
// way prohibited. For example, have a fork of https://github.com/Azure/autorest.go?
// Generate out and SDK with the same shape as our SDK, but using your patterns using
// this tool.
//
// Given that this code was developed as an internal tool, troubles with it that do not
// pertain to our usage of it may be slow to be fixed. Pull Requests are welcome though,
// feel free to contribute.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/marstr/collection"
)

var (
	targetFiles    collection.Enumerable
	packageVersion string
	outputBase     string
	logDirBase     string
	errLog         *log.Logger
	statusLog      *log.Logger
	dryRun         bool
)

// version should be set by the linker when compiling this program by providing the following arguments:
// -X main.version=<version>
//
// If installing the generator in your machine, that means running the following the command:
//   go install -ldflags "-X main.version=<version>"
//
// The reason this is not controlled in source, is to allow for the git commit SHA1 to be used as the
// version of the string. To retrieve the currently checked-out git commit SHA1, you can use the command:
//   git rev-parse HEAD
//
// If this value is set, it is recommended that it be the value of the SHA1 git commit identifier for the
// source that was used at build time.
var version string

func main() {
	start := time.Now()

	type generationTuple struct {
		fileName     string
		packageName  string
		outputFolder string
	}

	literateFiles := collection.Where(targetFiles, func(subject interface{}) bool {
		return strings.EqualFold(path.Base(subject.(string)), "README.md")
	})

	tuples := collection.SelectMany(literateFiles,
		// The following function compiles the regexp which finds Go related package tags in a literate file, and creates a collection.Unfolder.
		// This function has been declared this way so that the relatively expensive act of compiling a regexp is only done once.
		func() collection.Unfolder {
			const patternText = "```" + `\s+yaml\s+\$\(\s*tag\s*\)\s*==\s*'([\d\w\-\.]+)'\s*&&\s*\$\(go\)[\w\d\-\s:]*\s+output-folder:\s+\$\(go-sdk-folder\)([\w\d\-_\\/\.]+)\s+[\w\d\-\s:]*` + "```"

			goConfigPattern := regexp.MustCompile(patternText)

			// This function is a collection.Unfolder which takes a literate file as a path, and retrieves all configuration which applies to a package tag and Go.
			return func(subject interface{}) collection.Enumerator {
				results := make(chan interface{})

				go func() {
					defer close(results)

					targetContents, err := ioutil.ReadFile(subject.(string))
					if err != nil {
						errLog.Printf("Skipping %q because: %v", subject.(string), err)
						return
					}

					matches := goConfigPattern.FindAllStringSubmatch(string(targetContents), -1)

					if len(matches) == 0 {
						statusLog.Printf("Skipping %q because there were no package tags with go configuration found.", subject.(string))
					} else {
						for _, submatch := range matches {
							results <- generationTuple{
								fileName:     subject.(string),
								packageName:  normalizePath(submatch[1]),
								outputFolder: normalizePath(submatch[2]),
							}
						}
					}
				}()

				return results
			}
		}())

	if dryRun {
		for entry := range tuples.Enumerate(nil) {
			tuple := entry.(generationTuple)
			fmt.Printf("%q in %q to %q\n", tuple.packageName, tuple.fileName, tuple.outputFolder)
		}
	} else {
		var generatedCount, formattedCount, builtCount, vettedCount uint

		done := make(chan struct{})

		// Call AutoRest for each of the tags in each of the literate files, generating a Go package for calling that service.
		generated := tuples.Enumerate(done).ParallelSelect(func(subject interface{}) interface{} {
			tuple := subject.(generationTuple)
			args := []string{
				tuple.fileName,
				"--go",
				fmt.Sprintf("--go-sdk-folder='%s'", outputBase),
				"--verbose",
				"--tag=" + tuple.packageName,
				"--use=@microsoft.azure/autorest.go@v2.0.41",
			}

			if packageVersion != "" {
				args = append(args, fmt.Sprintf("--package-version='%s'", packageVersion))
				args = append(args, fmt.Sprintf("--user-agent='Azure-SDK-For-Go/%s services'", packageVersion))
			}

			logFileLoc := filepath.Join(logDirBase, tuple.outputFolder)
			err := os.MkdirAll(logFileLoc, os.ModePerm)
			if err != nil {
				errLog.Printf("Could not create log directory %q", logFileLoc)
				return nil
			}

			logFile, err := os.Create(filepath.Join(logFileLoc, "autorestLog.txt"))
			if err != nil {
				errLog.Printf("Could not create log file %q for AutoRest generating from: %q in %q", logFileLoc, tuple.packageName, tuple.fileName)
				return nil
			}

			commandText := new(bytes.Buffer)
			fmt.Fprint(commandText, "Executing Command: \"")
			fmt.Fprint(commandText, "autorest ")
			for _, a := range args {
				fmt.Fprint(commandText, a)
				fmt.Fprint(commandText, " ")
			}
			commandText.Truncate(commandText.Len() - 1)
			fmt.Fprint(commandText, `"`)

			fmt.Fprintln(logFile, commandText.String())

			genProc := exec.Command("autorest", args...)
			genProc.Stdout = logFile
			genProc.Stderr = logFile

			err = genProc.Run()
			if err != nil {
				fmt.Fprintln(logFile, "Autorest Exectution Error: ", err)
				return nil
			}
			generatedCount++
			return path.Join(outputBase, tuple.outputFolder)
		}).Where(isntNil)

		// Take all of the generated packages and reformat them to adhere to Go's guidelines.
		formatted := generated.Select(func(subject interface{}) (result interface{}) {
			err := exec.Command("gofmt", "-w", subject.(string)).Run()
			if err == nil {
				formattedCount++
				result = subject
			} else {
				errLog.Printf("Failed to format: %q", subject.(string))
			}
			return
		}).Where(isntNil)

		// Build all of the packages as a sanity check.
		built := formatted.Select(func(subject interface{}) (result interface{}) {
			pkgName := strings.TrimPrefix(trimGoPath(subject.(string)), "/src/")
			err := exec.Command("go", "build", pkgName).Run()
			if err == nil {
				builtCount++
				return pkgName
			}
			errLog.Printf("Failed to build: %q", pkgName)
			return nil
		}).Where(isntNil)

		vetted := built.Select(func(subject interface{}) interface{} {
			err := exec.Command("go", "vet", subject.(string)).Run()
			if err == nil {
				vettedCount++
				return subject
			}
			errLog.Printf("Failed to vet: %q", subject.(string))
			return nil
		}).Where(isntNil)

		// Turn the crank. This loop forces evaluation of the chain of enumerators that were built up
		// in the code above.
		for range vetted {
			// Intenionally Left Blank
		}

		fmt.Println("Execution Time: ", time.Now().Sub(start))
		fmt.Println("Generated: ", generatedCount)
		fmt.Println("Formatted: ", formattedCount)
		fmt.Println("Built: ", builtCount)
		fmt.Println("Vetted: ", vettedCount)
		close(done)
	}
}

func init() {
	var useRecursive, useStatus bool

	const automaticLogDirValue = "temp"
	const logDirUsage = "The root directory where all logs can be found. If the value `" + automaticLogDirValue + "` is supplied, a randomly named directory will be generated in the $TEMP directory."

	flag.BoolVar(&useRecursive, "r", false, "Recursively traverses the directories specified looking for literate files.")
	flag.StringVar(&outputBase, "o", getDefaultOutputBase(), "The root directory to use for the output of generated files. i.e. The value to be treated as the go-sdk-folder when AutoRest is called.")
	flag.BoolVar(&useStatus, "v", false, "Print status messages as generation takes place.")
	flag.BoolVar(&dryRun, "p", false, "Preview which packages would be generated instead of actaully calling autorest.")
	flag.StringVar(&packageVersion, "version", "", "The version that should be stamped on this SDK. This should be a semver.")
	flag.StringVar(&logDirBase, "l", getDefaultOutputBase(), logDirUsage)

	// Override the default usage message, printing to stderr as the default one would.
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "generator [options] [target Literate files]...")
		flag.PrintDefaults()
	}

	flag.Parse()

	statusWriter := ioutil.Discard
	if useStatus {
		statusWriter = os.Stdout
	}
	statusLog = log.New(statusWriter, "[STATUS] ", 0)
	errLog = log.New(os.Stderr, "[ERROR] ", 0)

	if logDirBase == automaticLogDirValue {
		var err error
		logDirBase, err = ioutil.TempDir("", "az-go-sdk-logs")
		logDirBase = normalizePath(logDirBase)
		if err == nil {
			statusLog.Print("Generation logs can be found at: ", logDirBase)
		} else {
			errLog.Print("Logging disabled. Could not create directory: ", logDirBase)
		}
	}

	targetFiles = collection.AsEnumerable(flag.Args())
	targetFiles = collection.SelectMany(targetFiles, func(subject interface{}) collection.Enumerator {
		cast, ok := subject.(string)

		if !ok {
			return collection.Empty.Enumerate(nil)
		}
		pathInfo, err := os.Stat(cast)
		if err != nil {
			return collection.Empty.Enumerate(nil)
		}

		if pathInfo.IsDir() {
			traverser := collection.Directory{
				Location: cast,
				Options:  collection.DirectoryOptionsExcludeDirectories,
			}

			if useRecursive {
				traverser.Options |= collection.DirectoryOptionsRecursive
			}

			return traverser.Enumerate(nil)
		}

		return collection.AsEnumerable(cast).Enumerate(nil)
	})

	targetFiles = collection.Select(targetFiles, func(subject interface{}) interface{} {
		return normalizePath(subject.(string))
	})
}

// goPath returns the value of the environement variable GOPATH at the beginning of this
// program's execution. The actual enviroment variable is only queried once, before any calls
// are made to this function.
var goPath = func() func() string {
	val := normalizePath(os.Getenv("GOPATH"))
	return func() string {
		return val
	}
}()

// getDefaultOutputBase returns the default location of the Azure-SDK-for-Go on your filesystem.
func getDefaultOutputBase() string {
	return normalizePath(path.Join(goPath(), "src", "github.com", "Azure", "azure-sdk-for-go"))
}

// trimGoPath operates like strings.TrimPrefix, where the prefix is always the value of GOPATH in the
// environment in which this program is being executed.
func trimGoPath(location string) string {
	return strings.TrimPrefix(normalizePath(location), goPath())
}

// normalizePath ensures that a path is expressed using forward slashes for sanity when working
// with Go Standard Library functions that seem to expect this.
func normalizePath(location string) (result string) {
	result = strings.Replace(location, `\`, "/", -1)
	return
}

// isntNil is a simple `collection.Predicate` which filters out `nil` objects from an Enumerator.
func isntNil(subject interface{}) bool {
	return subject != nil
}
<commit_msg>Make autorest.go version configurable. (#882)<commit_after>// Copyright 2017 Microsoft Corporation and contributors.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//
// See the License for the specific language governing permissions and
// limitations under the License.

// generator scours the github.com/Azure/azure-rest-api-specs repository in search for
// REST functionality that has Go configuration, then generates packages accordingly.
// This tool was developed with the intention that it be an internal tool for the
// Azure-SDK-for-Go team, but the usage of this package for public scenarios is in no
// way prohibited. For example, have a fork of https://github.com/Azure/autorest.go?
// Generate out and SDK with the same shape as our SDK, but using your patterns using
// this tool.
//
// Given that this code was developed as an internal tool, troubles with it that do not
// pertain to our usage of it may be slow to be fixed. Pull Requests are welcome though,
// feel free to contribute.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/marstr/collection"
)

const (
	defaultGenVer = "@microsoft.azure/autorest.go@v2.0.41"
)

var (
	targetFiles    collection.Enumerable
	packageVersion string
	outputBase     string
	logDirBase     string
	autorestVer    string
	errLog         *log.Logger
	statusLog      *log.Logger
	dryRun         bool
)

// version should be set by the linker when compiling this program by providing the following arguments:
// -X main.version=<version>
//
// If installing the generator in your machine, that means running the following the command:
//   go install -ldflags "-X main.version=<version>"
//
// The reason this is not controlled in source, is to allow for the git commit SHA1 to be used as the
// version of the string. To retrieve the currently checked-out git commit SHA1, you can use the command:
//   git rev-parse HEAD
//
// If this value is set, it is recommended that it be the value of the SHA1 git commit identifier for the
// source that was used at build time.
var version string

func main() {
	start := time.Now()

	type generationTuple struct {
		fileName     string
		packageName  string
		outputFolder string
	}

	literateFiles := collection.Where(targetFiles, func(subject interface{}) bool {
		return strings.EqualFold(path.Base(subject.(string)), "README.md")
	})

	tuples := collection.SelectMany(literateFiles,
		// The following function compiles the regexp which finds Go related package tags in a literate file, and creates a collection.Unfolder.
		// This function has been declared this way so that the relatively expensive act of compiling a regexp is only done once.
		func() collection.Unfolder {
			const patternText = "```" + `\s+yaml\s+\$\(\s*tag\s*\)\s*==\s*'([\d\w\-\.]+)'\s*&&\s*\$\(go\)[\w\d\-\s:]*\s+output-folder:\s+\$\(go-sdk-folder\)([\w\d\-_\\/\.]+)\s+[\w\d\-\s:]*` + "```"

			goConfigPattern := regexp.MustCompile(patternText)

			// This function is a collection.Unfolder which takes a literate file as a path, and retrieves all configuration which applies to a package tag and Go.
			return func(subject interface{}) collection.Enumerator {
				results := make(chan interface{})

				go func() {
					defer close(results)

					targetContents, err := ioutil.ReadFile(subject.(string))
					if err != nil {
						errLog.Printf("Skipping %q because: %v", subject.(string), err)
						return
					}

					matches := goConfigPattern.FindAllStringSubmatch(string(targetContents), -1)

					if len(matches) == 0 {
						statusLog.Printf("Skipping %q because there were no package tags with go configuration found.", subject.(string))
					} else {
						for _, submatch := range matches {
							results <- generationTuple{
								fileName:     subject.(string),
								packageName:  normalizePath(submatch[1]),
								outputFolder: normalizePath(submatch[2]),
							}
						}
					}
				}()

				return results
			}
		}())

	if dryRun {
		for entry := range tuples.Enumerate(nil) {
			tuple := entry.(generationTuple)
			fmt.Printf("%q in %q to %q\n", tuple.packageName, tuple.fileName, tuple.outputFolder)
		}
	} else {
		var generatedCount, formattedCount, builtCount, vettedCount uint

		done := make(chan struct{})

		// Call AutoRest for each of the tags in each of the literate files, generating a Go package for calling that service.
		generated := tuples.Enumerate(done).ParallelSelect(func(subject interface{}) interface{} {
			tuple := subject.(generationTuple)
			args := []string{
				tuple.fileName,
				"--go",
				fmt.Sprintf("--go-sdk-folder='%s'", outputBase),
				"--verbose",
				"--tag=" + tuple.packageName,
				"--use=" + autorestVer,
			}

			if packageVersion != "" {
				args = append(args, fmt.Sprintf("--package-version='%s'", packageVersion))
				args = append(args, fmt.Sprintf("--user-agent='Azure-SDK-For-Go/%s services'", packageVersion))
			}

			logFileLoc := filepath.Join(logDirBase, tuple.outputFolder)
			err := os.MkdirAll(logFileLoc, os.ModePerm)
			if err != nil {
				errLog.Printf("Could not create log directory %q", logFileLoc)
				return nil
			}

			logFile, err := os.Create(filepath.Join(logFileLoc, "autorestLog.txt"))
			if err != nil {
				errLog.Printf("Could not create log file %q for AutoRest generating from: %q in %q", logFileLoc, tuple.packageName, tuple.fileName)
				return nil
			}

			commandText := new(bytes.Buffer)
			fmt.Fprint(commandText, "Executing Command: \"")
			fmt.Fprint(commandText, "autorest ")
			for _, a := range args {
				fmt.Fprint(commandText, a)
				fmt.Fprint(commandText, " ")
			}
			commandText.Truncate(commandText.Len() - 1)
			fmt.Fprint(commandText, `"`)

			fmt.Fprintln(logFile, commandText.String())

			genProc := exec.Command("autorest", args...)
			genProc.Stdout = logFile
			genProc.Stderr = logFile

			err = genProc.Run()
			if err != nil {
				fmt.Fprintln(logFile, "Autorest Exectution Error: ", err)
				return nil
			}
			generatedCount++
			return path.Join(outputBase, tuple.outputFolder)
		}).Where(isntNil)

		// Take all of the generated packages and reformat them to adhere to Go's guidelines.
		formatted := generated.Select(func(subject interface{}) (result interface{}) {
			err := exec.Command("gofmt", "-w", subject.(string)).Run()
			if err == nil {
				formattedCount++
				result = subject
			} else {
				errLog.Printf("Failed to format: %q", subject.(string))
			}
			return
		}).Where(isntNil)

		// Build all of the packages as a sanity check.
		built := formatted.Select(func(subject interface{}) (result interface{}) {
			pkgName := strings.TrimPrefix(trimGoPath(subject.(string)), "/src/")
			err := exec.Command("go", "build", pkgName).Run()
			if err == nil {
				builtCount++
				return pkgName
			}
			errLog.Printf("Failed to build: %q", pkgName)
			return nil
		}).Where(isntNil)

		vetted := built.Select(func(subject interface{}) interface{} {
			err := exec.Command("go", "vet", subject.(string)).Run()
			if err == nil {
				vettedCount++
				return subject
			}
			errLog.Printf("Failed to vet: %q", subject.(string))
			return nil
		}).Where(isntNil)

		// Turn the crank. This loop forces evaluation of the chain of enumerators that were built up
		// in the code above.
		for range vetted {
			// Intenionally Left Blank
		}

		fmt.Println("Execution Time: ", time.Now().Sub(start))
		fmt.Println("Generated: ", generatedCount)
		fmt.Println("Formatted: ", formattedCount)
		fmt.Println("Built: ", builtCount)
		fmt.Println("Vetted: ", vettedCount)
		close(done)
	}
}

func init() {
	var useRecursive, useStatus bool

	const automaticLogDirValue = "temp"
	const logDirUsage = "The root directory where all logs can be found. If the value `" + automaticLogDirValue + "` is supplied, a randomly named directory will be generated in the $TEMP directory."

	flag.BoolVar(&useRecursive, "r", false, "Recursively traverses the directories specified looking for literate files.")
	flag.StringVar(&outputBase, "o", getDefaultOutputBase(), "The root directory to use for the output of generated files. i.e. The value to be treated as the go-sdk-folder when AutoRest is called.")
	flag.BoolVar(&useStatus, "v", false, "Print status messages as generation takes place.")
	flag.BoolVar(&dryRun, "p", false, "Preview which packages would be generated instead of actaully calling autorest.")
	flag.StringVar(&packageVersion, "version", "", "The version that should be stamped on this SDK. This should be a semver.")
	flag.StringVar(&logDirBase, "l", getDefaultOutputBase(), logDirUsage)
	flag.StringVar(&autorestVer, "a", defaultGenVer, "The version of the AutoRest Go code generator to use, defaults to `"+defaultGenVer+"`.")

	// Override the default usage message, printing to stderr as the default one would.
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "generator [options] [target Literate files]...")
		flag.PrintDefaults()
	}

	flag.Parse()

	statusWriter := ioutil.Discard
	if useStatus {
		statusWriter = os.Stdout
	}
	statusLog = log.New(statusWriter, "[STATUS] ", 0)
	errLog = log.New(os.Stderr, "[ERROR] ", 0)

	if logDirBase == automaticLogDirValue {
		var err error
		logDirBase, err = ioutil.TempDir("", "az-go-sdk-logs")
		logDirBase = normalizePath(logDirBase)
		if err == nil {
			statusLog.Print("Generation logs can be found at: ", logDirBase)
		} else {
			errLog.Print("Logging disabled. Could not create directory: ", logDirBase)
		}
	}

	targetFiles = collection.AsEnumerable(flag.Args())
	targetFiles = collection.SelectMany(targetFiles, func(subject interface{}) collection.Enumerator {
		cast, ok := subject.(string)

		if !ok {
			return collection.Empty.Enumerate(nil)
		}
		pathInfo, err := os.Stat(cast)
		if err != nil {
			return collection.Empty.Enumerate(nil)
		}

		if pathInfo.IsDir() {
			traverser := collection.Directory{
				Location: cast,
				Options:  collection.DirectoryOptionsExcludeDirectories,
			}

			if useRecursive {
				traverser.Options |= collection.DirectoryOptionsRecursive
			}

			return traverser.Enumerate(nil)
		}

		return collection.AsEnumerable(cast).Enumerate(nil)
	})

	targetFiles = collection.Select(targetFiles, func(subject interface{}) interface{} {
		return normalizePath(subject.(string))
	})
}

// goPath returns the value of the environement variable GOPATH at the beginning of this
// program's execution. The actual enviroment variable is only queried once, before any calls
// are made to this function.
var goPath = func() func() string {
	val := normalizePath(os.Getenv("GOPATH"))
	return func() string {
		return val
	}
}()

// getDefaultOutputBase returns the default location of the Azure-SDK-for-Go on your filesystem.
func getDefaultOutputBase() string {
	return normalizePath(path.Join(goPath(), "src", "github.com", "Azure", "azure-sdk-for-go"))
}

// trimGoPath operates like strings.TrimPrefix, where the prefix is always the value of GOPATH in the
// environment in which this program is being executed.
func trimGoPath(location string) string {
	return strings.TrimPrefix(normalizePath(location), goPath())
}

// normalizePath ensures that a path is expressed using forward slashes for sanity when working
// with Go Standard Library functions that seem to expect this.
func normalizePath(location string) (result string) {
	result = strings.Replace(location, `\`, "/", -1)
	return
}

// isntNil is a simple `collection.Predicate` which filters out `nil` objects from an Enumerator.
func isntNil(subject interface{}) bool {
	return subject != nil
}
<|endoftext|>


<commit_before>package wipro

import (
	"fmt"
	"io"

	"h12.me/gengo"
)

func (t GoTypes) GoFuncs(w io.Writer) {
	fpl(w, "package proto")
	fpl(w, "import (")
	fpl(w, `"hash/crc32"`)
	fpl(w, ")")
	for _, decl := range t.TypeDecls {
		genMarshalFunc(w, decl)
		fpl(w, "")
		genUnmarshalFunc(w, decl)
		fpl(w, "")
	}
}

func genMarshalFunc(w io.Writer, decl *gengo.TypeDecl) {
	t := &decl.Type
	if t.Kind == gengo.IdentKind && t.Ident == "M" {
		return
	}
	fpl(w, "func (t *%s) Marshal(w *Writer) {", decl.Name)
	switch t.Kind {
	case gengo.StructKind:
		if f0 := t.Fields[0]; len(t.Fields) == 2 &&
			(f0.Name == "Size" || f0.Name == "CRC") {
			fpl(w, "offset := len(w.B)")
			marshalField(w, t.Fields[0])
			fpl(w, "start := len(w.B)")
			marshalField(w, t.Fields[1])
			switch f0.Name {
			case "Size":
				fpl(w, "w.SetInt32(offset, int32(len(w.B)-start))")
			case "CRC":
				fpl(w, "w.SetUint32(offset, crc32.ChecksumIEEE(w.B[start:]))")
			}
		} else {
			for _, f := range t.Fields {
				marshalField(w, f)
			}
		}
	case gengo.ArrayKind:
		marshalValue(w, "(*t)", t)
	case gengo.IdentKind:
		marshalValue(w, "(*t)", t)
	default:
		fpl(w, "// type %s, %v", decl.Name, decl.Type)
	}
	fpl(w, "}")
}

func genUnmarshalFunc(w io.Writer, decl *gengo.TypeDecl) {
	t := &decl.Type
	if t.Kind == gengo.IdentKind && t.Ident == "M" {
		return
	}
	fpl(w, "func (t *%s) Unmarshal(r *Reader) {", decl.Name)
	switch t.Kind {
	case gengo.StructKind:
		if f0 := t.Fields[0]; len(t.Fields) == 2 &&
			(f0.Name == "Size" || f0.Name == "CRC") {
			unmarshalField(w, t.Fields[0])
			fpl(w, "start := r.Offset")
			unmarshalField(w, t.Fields[1])
			switch f0.Name {
			case "Size":
				fpl(w, "if r.Err == nil && int(t.Size) != r.Offset-start {")
				fpl(w, `r.Err = newError("size mismatch, expect %%d, got %%d", int(t.Size), r.Offset-start)`)
				fpl(w, "}")
			case "CRC":
				fpl(w, "if r.Err == nil && t.CRC != crc32.ChecksumIEEE(r.B[start:r.Offset]) {")
				fpl(w, `r.Err = newError("CRC mismatch")`)
				fpl(w, "}")
			}
		} else {
			for _, f := range t.Fields {
				unmarshalField(w, f)
			}
		}
	case gengo.ArrayKind:
		unmarshalValue(w, "(*t)", t, t.Ident)
	case gengo.IdentKind:
		unmarshalValue(w, "(*t)", t, decl.Name)
	default:
		fpl(w, "// type %s, %v", decl.Name, decl.Type)
	}
	fpl(w, "}")
}

func marshalField(w io.Writer, f *gengo.Field) {
	fName := "t." + f.Name
	if f.Name == "" {
		fName = "t." + f.Type.Ident
	}
	marshalValue(w, fName, &f.Type)
}

func unmarshalField(w io.Writer, f *gengo.Field) {
	fName := "t." + f.Name
	if f.Name == "" {
		fName = "t." + f.Type.Ident
	}
	unmarshalValue(w, fName, &f.Type, f.Type.Ident)
}

func marshalValue(w io.Writer, name string, typ *gengo.Type) {
	switch typ.Kind {
	case gengo.IdentKind:
		switch typ.Ident {
		case "int64":
			marshalInt(w, name, 64)
		case "int32":
			marshalInt(w, name, 32)
		case "int16":
			marshalInt(w, name, 16)
		case "int8":
			marshalInt(w, name, 8)
		case "uint64":
			marshalUint(w, name, 64)
		case "uint32":
			marshalUint(w, name, 32)
		case "uint16":
			marshalUint(w, name, 16)
		case "uint8":
			marshalUint(w, name, 8)
		case "string":
			fpl(w, "w.WriteString(string(%s))", name)
		case "[]byte":
			fpl(w, "w.WriteBytes(%s)", name)
		default:
			marshalMarshaler(w, name)
		}
	case gengo.ArrayKind:
		switch typ.Get("array_prefix").(string) {
		case "length":
			marshalLengthArray(w, name, typ)
		case "size":
			marshalSizeArray(w, name, typ)
		}
	default:
		fpl(w, "// value %s %v", name, typ.Kind)
	}
}

func unmarshalValue(w io.Writer, name string, typ *gengo.Type, declType string) {
	switch typ.Kind {
	case gengo.IdentKind:
		switch typ.Ident {
		case "int64":
			unmarshalInt(w, name, "b", 64)
		case "int32":
			unmarshalInt(w, name, "b", 32)
		case "int16":
			unmarshalInt(w, name, "b", 16)
		case "int8":
			unmarshalInt(w, name, "b", 8)
		case "uint64":
			unmarshalUint(w, name, "b", 64)
		case "uint32":
			unmarshalUint(w, name, "b", 32)
		case "uint16":
			unmarshalUint(w, name, "b", 16)
		case "uint8":
			unmarshalUint(w, name, "b", 8)
		case "string":
			fpl(w, "%s = %s(r.ReadString())", name, declType)
		case "[]byte":
			fpl(w, "%s = r.ReadBytes()", name)
		default:
			unmarshalUnmarshaler(w, name, typ.Ident)
		}
	case gengo.ArrayKind:
		switch typ.Get("array_prefix").(string) {
		case "length":
			unmarshalLengthArray(w, name, typ)
		case "size":
			unmarshalSizeArray(w, name, typ)
		}
	default:
		fpl(w, "// value %s %v", name, typ.Kind)
	}
}

func marshalLengthArray(w io.Writer, name string, typ *gengo.Type) {
	marshalInt(w, fmt.Sprintf("int32(len(%s))", name), 32)
	fpl(w, "for i := range %s {", name)
	switch typ.Ident {
	case "int8", "int16", "int32", "int64", "string":
		marshalValue(w, name+"[i]", &gengo.Type{Kind: gengo.IdentKind, Ident: typ.Ident})
	default:
		marshalMarshaler(w, name+"[i]")
	}
	fpl(w, "}")
}

func unmarshalLengthArray(w io.Writer, name string, typ *gengo.Type) {
	fpl(w, "%s = make([]%s, int(r.ReadInt32()))", name, typ.Ident)
	fpl(w, "for i := range %s {", name)
	switch typ.Ident {
	case "int8", "int16", "int32", "int64", "string":
		unmarshalValue(w, name+"[i]", &gengo.Type{Kind: gengo.IdentKind, Ident: typ.Ident}, typ.Ident)
	default:
		unmarshalUnmarshaler(w, name+"[i]", typ.Ident)
	}
	fpl(w, "}")
}

func marshalSizeArray(w io.Writer, name string, typ *gengo.Type) {
	fpl(w, "offset := len(w.B)")
	fpl(w, "w.WriteInt32(0)")
	fpl(w, "start := len(w.B)")
	fpl(w, "for i := range %s {", name)
	switch typ.Ident {
	case "int8", "int16", "int32", "int64", "string":
		marshalValue(w, name+"[i]", &gengo.Type{Kind: gengo.IdentKind, Ident: typ.Ident})
	default:
		marshalMarshaler(w, name+"[i]")
	}
	fpl(w, "}")
	fpl(w, "w.SetInt32(offset, int32(len(w.B)-start))")
}

func unmarshalSizeArray(w io.Writer, name string, typ *gengo.Type) {
	fpl(w, "size := int(r.ReadInt32())")
	fpl(w, "start := r.Offset")
	fpl(w, "for r.Offset-start < size {")
	fpl(w, "var m %s", typ.Ident)
	switch typ.Ident {
	case "int8", "int16", "int32", "int64", "string":
		unmarshalValue(w, "m", &gengo.Type{Kind: gengo.IdentKind, Ident: typ.Ident}, typ.Ident)
	default:
		unmarshalUnmarshaler(w, "m", typ.Ident)
	}
	fpl(w, "if r.Err != nil {")
	fpl(w, "return")
	fpl(w, "}")
	fpl(w, "*t = append(*t, m)")
	fpl(w, "}")
}

func marshalMarshaler(w io.Writer, marshaler string) {
	fpl(w, "%s.Marshal(w)", marshaler)
}

func unmarshalUnmarshaler(w io.Writer, unmarshaler, typ string) {
	fpl(w, "%s.Unmarshal(r)", unmarshaler)
}

func marshalInt(w io.Writer, name string, bit int) {
	fpl(w, "w.WriteInt%d(%s)", bit, name)
}

func unmarshalInt(w io.Writer, name string, bufName string, bit int) {
	fpl(w, "%s=r.ReadInt%d()", name, bit)
}

func marshalUint(w io.Writer, name string, bit int) {
	fpl(w, "w.WriteUint%d(%s)", bit, name)
}

func unmarshalUint(w io.Writer, name string, bufName string, bit int) {
	fpl(w, "%s=r.ReadUint%d()", name, bit)
}
<commit_msg>allow CRC and Sized struct have more than two fields.<commit_after>package wipro

import (
	"fmt"
	"io"

	"h12.me/gengo"
)

func (t GoTypes) GoFuncs(w io.Writer) {
	fpl(w, "package proto")
	fpl(w, "import (")
	fpl(w, `"hash/crc32"`)
	fpl(w, ")")
	for _, decl := range t.TypeDecls {
		genMarshalFunc(w, decl)
		fpl(w, "")
		genUnmarshalFunc(w, decl)
		fpl(w, "")
	}
}

func genMarshalFunc(w io.Writer, decl *gengo.TypeDecl) {
	t := &decl.Type
	if t.Kind == gengo.IdentKind && t.Ident == "M" {
		return
	}
	fpl(w, "func (t *%s) Marshal(w *Writer) {", decl.Name)
	switch t.Kind {
	case gengo.StructKind:
		if f0 := t.Fields[0]; f0.Name == "Size" || f0.Name == "CRC" {
			fpl(w, "offset := len(w.B)")
			marshalField(w, t.Fields[0])
			fpl(w, "start := len(w.B)")
			for _, field := range t.Fields[1:] {
				marshalField(w, field)
			}
			switch f0.Name {
			case "Size":
				fpl(w, "w.SetInt32(offset, int32(len(w.B)-start))")
			case "CRC":
				fpl(w, "w.SetUint32(offset, crc32.ChecksumIEEE(w.B[start:]))")
			}
		} else {
			for _, f := range t.Fields {
				marshalField(w, f)
			}
		}
	case gengo.ArrayKind:
		marshalValue(w, "(*t)", t)
	case gengo.IdentKind:
		marshalValue(w, "(*t)", t)
	default:
		fpl(w, "// type %s, %v", decl.Name, decl.Type)
	}
	fpl(w, "}")
}

func genUnmarshalFunc(w io.Writer, decl *gengo.TypeDecl) {
	t := &decl.Type
	if t.Kind == gengo.IdentKind && t.Ident == "M" {
		return
	}
	fpl(w, "func (t *%s) Unmarshal(r *Reader) {", decl.Name)
	switch t.Kind {
	case gengo.StructKind:
		if f0 := t.Fields[0]; f0.Name == "Size" || f0.Name == "CRC" {
			unmarshalField(w, t.Fields[0])
			fpl(w, "start := r.Offset")
			for _, field := range t.Fields[1:] {
				unmarshalField(w, field)
			}
			switch f0.Name {
			case "Size":
				fpl(w, "if r.Err == nil && int(t.Size) != r.Offset-start {")
				fpl(w, `r.Err = ErrSizeMismatch`)
				fpl(w, "}")
			case "CRC":
				fpl(w, "if r.Err == nil && t.CRC != crc32.ChecksumIEEE(r.B[start:r.Offset]) {")
				fpl(w, `r.Err = ErrCRCMismatch`)
				fpl(w, "}")
			}
		} else {
			for _, f := range t.Fields {
				unmarshalField(w, f)
			}
		}
	case gengo.ArrayKind:
		unmarshalValue(w, "(*t)", t, t.Ident)
	case gengo.IdentKind:
		unmarshalValue(w, "(*t)", t, decl.Name)
	default:
		fpl(w, "// type %s, %v", decl.Name, decl.Type)
	}
	fpl(w, "}")
}

func marshalField(w io.Writer, f *gengo.Field) {
	fName := "t." + f.Name
	if f.Name == "" {
		fName = "t." + f.Type.Ident
	}
	marshalValue(w, fName, &f.Type)
}

func unmarshalField(w io.Writer, f *gengo.Field) {
	fName := "t." + f.Name
	if f.Name == "" {
		fName = "t." + f.Type.Ident
	}
	unmarshalValue(w, fName, &f.Type, f.Type.Ident)
}

func marshalValue(w io.Writer, name string, typ *gengo.Type) {
	switch typ.Kind {
	case gengo.IdentKind:
		switch typ.Ident {
		case "int64":
			marshalInt(w, name, 64)
		case "int32":
			marshalInt(w, name, 32)
		case "int16":
			marshalInt(w, name, 16)
		case "int8":
			marshalInt(w, name, 8)
		case "uint64":
			marshalUint(w, name, 64)
		case "uint32":
			marshalUint(w, name, 32)
		case "uint16":
			marshalUint(w, name, 16)
		case "uint8":
			marshalUint(w, name, 8)
		case "string":
			fpl(w, "w.WriteString(string(%s))", name)
		case "[]byte":
			fpl(w, "w.WriteBytes(%s)", name)
		default:
			marshalMarshaler(w, name)
		}
	case gengo.ArrayKind:
		switch typ.Get("array_prefix").(string) {
		case "length":
			marshalLengthArray(w, name, typ)
		case "size":
			marshalSizeArray(w, name, typ)
		}
	default:
		fpl(w, "// value %s %v", name, typ.Kind)
	}
}

func unmarshalValue(w io.Writer, name string, typ *gengo.Type, declType string) {
	switch typ.Kind {
	case gengo.IdentKind:
		switch typ.Ident {
		case "int64":
			unmarshalInt(w, name, "b", 64)
		case "int32":
			unmarshalInt(w, name, "b", 32)
		case "int16":
			unmarshalInt(w, name, "b", 16)
		case "int8":
			unmarshalInt(w, name, "b", 8)
		case "uint64":
			unmarshalUint(w, name, "b", 64)
		case "uint32":
			unmarshalUint(w, name, "b", 32)
		case "uint16":
			unmarshalUint(w, name, "b", 16)
		case "uint8":
			unmarshalUint(w, name, "b", 8)
		case "string":
			fpl(w, "%s = %s(r.ReadString())", name, declType)
		case "[]byte":
			fpl(w, "%s = r.ReadBytes()", name)
		default:
			unmarshalUnmarshaler(w, name, typ.Ident)
		}
	case gengo.ArrayKind:
		switch typ.Get("array_prefix").(string) {
		case "length":
			unmarshalLengthArray(w, name, typ)
		case "size":
			unmarshalSizeArray(w, name, typ)
		}
	default:
		fpl(w, "// value %s %v", name, typ.Kind)
	}
}

func marshalLengthArray(w io.Writer, name string, typ *gengo.Type) {
	marshalInt(w, fmt.Sprintf("int32(len(%s))", name), 32)
	fpl(w, "for i := range %s {", name)
	switch typ.Ident {
	case "int8", "int16", "int32", "int64", "string":
		marshalValue(w, name+"[i]", &gengo.Type{Kind: gengo.IdentKind, Ident: typ.Ident})
	default:
		marshalMarshaler(w, name+"[i]")
	}
	fpl(w, "}")
}

func unmarshalLengthArray(w io.Writer, name string, typ *gengo.Type) {
	fpl(w, "%s = make([]%s, int(r.ReadInt32()))", name, typ.Ident)
	fpl(w, "for i := range %s {", name)
	switch typ.Ident {
	case "int8", "int16", "int32", "int64", "string":
		unmarshalValue(w, name+"[i]", &gengo.Type{Kind: gengo.IdentKind, Ident: typ.Ident}, typ.Ident)
	default:
		unmarshalUnmarshaler(w, name+"[i]", typ.Ident)
	}
	fpl(w, "}")
}

func marshalSizeArray(w io.Writer, name string, typ *gengo.Type) {
	fpl(w, "offset := len(w.B)")
	fpl(w, "w.WriteInt32(0)")
	fpl(w, "start := len(w.B)")
	fpl(w, "for i := range %s {", name)
	switch typ.Ident {
	case "int8", "int16", "int32", "int64", "string":
		marshalValue(w, name+"[i]", &gengo.Type{Kind: gengo.IdentKind, Ident: typ.Ident})
	default:
		marshalMarshaler(w, name+"[i]")
	}
	fpl(w, "}")
	fpl(w, "w.SetInt32(offset, int32(len(w.B)-start))")
}

func unmarshalSizeArray(w io.Writer, name string, typ *gengo.Type) {
	fpl(w, "size := int(r.ReadInt32())")
	fpl(w, "start := r.Offset")
	fpl(w, "for r.Offset-start < size {")
	fpl(w, "var m %s", typ.Ident)
	switch typ.Ident {
	case "int8", "int16", "int32", "int64", "string":
		unmarshalValue(w, "m", &gengo.Type{Kind: gengo.IdentKind, Ident: typ.Ident}, typ.Ident)
	default:
		unmarshalUnmarshaler(w, "m", typ.Ident)
	}
	fpl(w, "if r.Err != nil {")
	fpl(w, "return")
	fpl(w, "}")
	fpl(w, "*t = append(*t, m)")
	fpl(w, "}")
}

func marshalMarshaler(w io.Writer, marshaler string) {
	fpl(w, "%s.Marshal(w)", marshaler)
}

func unmarshalUnmarshaler(w io.Writer, unmarshaler, typ string) {
	fpl(w, "%s.Unmarshal(r)", unmarshaler)
}

func marshalInt(w io.Writer, name string, bit int) {
	fpl(w, "w.WriteInt%d(%s)", bit, name)
}

func unmarshalInt(w io.Writer, name string, bufName string, bit int) {
	fpl(w, "%s=r.ReadInt%d()", name, bit)
}

func marshalUint(w io.Writer, name string, bit int) {
	fpl(w, "w.WriteUint%d(%s)", bit, name)
}

func unmarshalUint(w io.Writer, name string, bufName string, bit int) {
	fpl(w, "%s=r.ReadUint%d()", name, bit)
}
<|endoftext|>


<commit_before>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"

	"github.com/google/syzkaller/pkg/csource"
	"github.com/google/syzkaller/prog"
	_ "github.com/google/syzkaller/sys"
)

var (
	flagOS         = flag.String("os", runtime.GOOS, "target os")
	flagArch       = flag.String("arch", runtime.GOARCH, "target arch")
	flagThreaded   = flag.Bool("threaded", false, "create threaded program")
	flagCollide    = flag.Bool("collide", false, "create collide program")
	flagRepeat     = flag.Bool("repeat", false, "repeat program infinitely or not")
	flagProcs      = flag.Int("procs", 1, "number of parallel processes")
	flagSandbox    = flag.String("sandbox", "", "sandbox to use (none, setuid, namespace)")
	flagProg       = flag.String("prog", "", "file with program to convert (required)")
	flagFaultCall  = flag.Int("fault_call", -1, "inject fault into this call (0-based)")
	flagFaultNth   = flag.Int("fault_nth", 0, "inject fault on n-th operation (0-based)")
	flagEnableTun  = flag.Bool("tun", false, "set up TUN/TAP interface")
	flagUseTmpDir  = flag.Bool("tmpdir", false, "create a temporary dir and execute inside it")
	flagHandleSegv = flag.Bool("segv", false, "catch and ignore SIGSEGV")
	flagWaitRepeat = flag.Bool("waitrepeat", false, "wait for each repeat attempt")
	flagDebug      = flag.Bool("debug", false, "generate debug printfs")
)

func main() {
	flag.Parse()
	if *flagProg == "" {
		flag.PrintDefaults()
		os.Exit(1)
	}
	target, err := prog.GetTarget(*flagOS, *flagArch)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v", err)
		os.Exit(1)
	}
	data, err := ioutil.ReadFile(*flagProg)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to read prog file: %v\n", err)
		os.Exit(1)
	}
	p, err := target.Deserialize(data)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to deserialize the program: %v\n", err)
		os.Exit(1)
	}
	opts := csource.Options{
		Threaded:   *flagThreaded,
		Collide:    *flagCollide,
		Repeat:     *flagRepeat,
		Procs:      *flagProcs,
		Sandbox:    *flagSandbox,
		Fault:      *flagFaultCall >= 0,
		FaultCall:  *flagFaultCall,
		FaultNth:   *flagFaultNth,
		EnableTun:  *flagEnableTun,
		UseTmpDir:  *flagUseTmpDir,
		HandleSegv: *flagHandleSegv,
		WaitRepeat: *flagWaitRepeat,
		Debug:      *flagDebug,
		Repro:      false,
	}
	src, err := csource.Write(p, opts)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to generate C source: %v\n", err)
		os.Exit(1)
	}
	if formatted, err := csource.Format(src); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
	} else {
		src = formatted
	}
	os.Stdout.Write(src)
}
<commit_msg>tools/syz-prog2c: add missing csource option<commit_after>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"

	"github.com/google/syzkaller/pkg/csource"
	"github.com/google/syzkaller/prog"
	_ "github.com/google/syzkaller/sys"
)

var (
	flagOS         = flag.String("os", runtime.GOOS, "target os")
	flagArch       = flag.String("arch", runtime.GOARCH, "target arch")
	flagThreaded   = flag.Bool("threaded", false, "create threaded program")
	flagCollide    = flag.Bool("collide", false, "create collide program")
	flagRepeat     = flag.Bool("repeat", false, "repeat program infinitely or not")
	flagProcs      = flag.Int("procs", 1, "number of parallel processes")
	flagSandbox    = flag.String("sandbox", "", "sandbox to use (none, setuid, namespace)")
	flagProg       = flag.String("prog", "", "file with program to convert (required)")
	flagFaultCall  = flag.Int("fault_call", -1, "inject fault into this call (0-based)")
	flagFaultNth   = flag.Int("fault_nth", 0, "inject fault on n-th operation (0-based)")
	flagEnableTun  = flag.Bool("tun", false, "set up TUN/TAP interface")
	flagUseTmpDir  = flag.Bool("tmpdir", false, "create a temporary dir and execute inside it")
	flagHandleSegv = flag.Bool("segv", false, "catch and ignore SIGSEGV")
	flagWaitRepeat = flag.Bool("waitrepeat", false, "wait for each repeat attempt")
	flagCgroups = flag.Bool("cgroups", false, "enable cgroups support")
	flagDebug      = flag.Bool("debug", false, "generate debug printfs")
)

func main() {
	flag.Parse()
	if *flagProg == "" {
		flag.PrintDefaults()
		os.Exit(1)
	}
	target, err := prog.GetTarget(*flagOS, *flagArch)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v", err)
		os.Exit(1)
	}
	data, err := ioutil.ReadFile(*flagProg)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to read prog file: %v\n", err)
		os.Exit(1)
	}
	p, err := target.Deserialize(data)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to deserialize the program: %v\n", err)
		os.Exit(1)
	}
	opts := csource.Options{
		Threaded:   *flagThreaded,
		Collide:    *flagCollide,
		Repeat:     *flagRepeat,
		Procs:      *flagProcs,
		Sandbox:    *flagSandbox,
		Fault:      *flagFaultCall >= 0,
		FaultCall:  *flagFaultCall,
		FaultNth:   *flagFaultNth,
		EnableTun:  *flagEnableTun,
		UseTmpDir:  *flagUseTmpDir,
		EnableCgroups: *flagCgroups,
		HandleSegv: *flagHandleSegv,
		WaitRepeat: *flagWaitRepeat,
		Debug:      *flagDebug,
		Repro:      false,
	}
	src, err := csource.Write(p, opts)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to generate C source: %v\n", err)
		os.Exit(1)
	}
	if formatted, err := csource.Format(src); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
	} else {
		src = formatted
	}
	os.Stdout.Write(src)
}
<|endoftext|>


<commit_before>package match

import (
	"reflect"
	"testing"
)

// Matcher is a helper object for chaining matcher conditions.
type Matcher struct {
	t     *testing.T
	value interface{}
}

// Fatal causes the unit test to immediately fail is one of the previous
// conditions has fail.
// Note: This will only affect previous conditions, not subsequent ones. You
// can build complex conditions this way.
// Ex: match.IsNotNil(t, value).Fatal().Equals(other)
func (m *Matcher) Fatal() *Matcher {
	if m.t.Failed() {
		m.t.FailNow()
	}
	return m
}

// IsNil asserts that the matched value is nil.
func (m *Matcher) IsNil() *Matcher {
	if m.value != nil {
		m.t.Errorf("expected == nil, got %v", m.value)
	}
	return m
}

// IsNotNil asserts that the matched value is not nil.
func (m *Matcher) IsNotNil() *Matcher {
	if m.value == nil {
		m.t.Error("expected != nil, got nil")
	}
	return m
}

// Equals asserts that the matched value is equal to value.
func (m *Matcher) Equals(value interface{}) *Matcher {
	if !reflect.DeepEqual(m.value, value) {
		m.t.Errorf("expected == %v, got %v", value, m.value)
	}
	return m
}

// NotEquals asserts that the match value is not equal to value.
func (m *Matcher) NotEquals(value interface{}) *Matcher {
	if reflect.DeepEqual(m.value, value) {
		m.t.Errorf("expected != %v, got %v", value, m.value)
	}
	return m
}

// LessThan asserts that the matched value is less than value.
// Note: This function will only work on numeric values.
// If a non-numeric value is passed, a fatal error will be thrown.
func (m *Matcher) LessThan(value interface{}) *Matcher {
	ok := true
	rv := reflect.ValueOf(m.value)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		ok = rv.Int() < reflect.ValueOf(value).Int()
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		ok = rv.Uint() < uint64(reflect.ValueOf(value).Int())
	case reflect.Float32, reflect.Float64:
		ok = rv.Float() < reflect.ValueOf(value).Float()
	default:
		m.t.Fatalf("expected numeric value, got %v", m.value)
	}
	if !ok {
		m.t.Errorf("expected < %v, got %v", value, m.value)
	}
	return m
}

// GreaterThan asserts that the matched value is greater then value.
// Note: This function will only work on numeric values.
// If a non-numeric value is passed, a fatal error will be thrown.
func (m *Matcher) GreaterThan(value interface{}) *Matcher {
	ok := true
	rv := reflect.ValueOf(m.value)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		ok = rv.Int() > reflect.ValueOf(value).Int()
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		ok = rv.Uint() > uint64(reflect.ValueOf(value).Int())
	case reflect.Float32, reflect.Float64:
		ok = rv.Float() > reflect.ValueOf(value).Float()
	default:
		m.t.Fatalf("expected numeric value, got %v", m.value)
	}
	if !ok {
		m.t.Errorf("expected < %v, got %v", value, m.value)
	}
	return m
}

// Contains asserts that the matched value contains value.
// If used on an array, asserts that the array contains value.
// If used on a map, asserts that the map contains a key with value.
// If used on a string, asserts that the string contains value.
func (m *Matcher) Contains(value interface{}) *Matcher {
	// TODO: Implement Matcher.Contains
	return m
}

// Matches asserts that the matched value matches the regex in pattern.
// Note: If the regex cannot be compiled, a fatal error will be thrown.
// If used on a non-string value, a fatal error will be thrown.
func (m *Matcher) Matches(pattern string) *Matcher {
	// TODO: Implement Matcher.Matches
	return m
}

// KindOf asserts that the matched value of the kind.
func (m *Matcher) KindOf(t reflect.Kind) *Matcher {
	// TODO: Implement Matcher.KindOf
	return m
}
<commit_msg>Implement Matches matcher<commit_after>package match

import (
	"reflect"
	"regexp"
	"testing"
)

// Matcher is a helper object for chaining matcher conditions.
type Matcher struct {
	t     *testing.T
	value interface{}
}

// Fatal causes the unit test to immediately fail is one of the previous
// conditions has fail.
// Note: This will only affect previous conditions, not subsequent ones. You
// can build complex conditions this way.
// Ex: match.IsNotNil(t, value).Fatal().Equals(other)
func (m *Matcher) Fatal() *Matcher {
	if m.t.Failed() {
		m.t.FailNow()
	}
	return m
}

// IsNil asserts that the matched value is nil.
func (m *Matcher) IsNil() *Matcher {
	if m.value != nil {
		m.t.Errorf("expected == nil, got %v", m.value)
	}
	return m
}

// IsNotNil asserts that the matched value is not nil.
func (m *Matcher) IsNotNil() *Matcher {
	if m.value == nil {
		m.t.Error("expected != nil, got nil")
	}
	return m
}

// Equals asserts that the matched value is equal to value.
func (m *Matcher) Equals(value interface{}) *Matcher {
	if !reflect.DeepEqual(m.value, value) {
		m.t.Errorf("expected == %v, got %v", value, m.value)
	}
	return m
}

// NotEquals asserts that the match value is not equal to value.
func (m *Matcher) NotEquals(value interface{}) *Matcher {
	if reflect.DeepEqual(m.value, value) {
		m.t.Errorf("expected != %v, got %v", value, m.value)
	}
	return m
}

// LessThan asserts that the matched value is less than value.
// Note: This function will only work on numeric values.
// If a non-numeric value is passed, a fatal error will be thrown.
func (m *Matcher) LessThan(value interface{}) *Matcher {
	ok := true
	rv := reflect.ValueOf(m.value)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		ok = rv.Int() < reflect.ValueOf(value).Int()
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		ok = rv.Uint() < uint64(reflect.ValueOf(value).Int())
	case reflect.Float32, reflect.Float64:
		ok = rv.Float() < reflect.ValueOf(value).Float()
	default:
		m.t.Fatalf("expected numeric value, got %v", m.value)
	}
	if !ok {
		m.t.Errorf("expected < %v, got %v", value, m.value)
	}
	return m
}

// GreaterThan asserts that the matched value is greater then value.
// Note: This function will only work on numeric values.
// If a non-numeric value is passed, a fatal error will be thrown.
func (m *Matcher) GreaterThan(value interface{}) *Matcher {
	ok := true
	rv := reflect.ValueOf(m.value)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		ok = rv.Int() > reflect.ValueOf(value).Int()
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		ok = rv.Uint() > uint64(reflect.ValueOf(value).Int())
	case reflect.Float32, reflect.Float64:
		ok = rv.Float() > reflect.ValueOf(value).Float()
	default:
		m.t.Fatalf("expected numeric value, got %v", m.value)
	}
	if !ok {
		m.t.Errorf("expected < %v, got %v", value, m.value)
	}
	return m
}

// Contains asserts that the matched value contains value.
// If used on an array, asserts that the array contains value.
// If used on a map, asserts that the map contains a key with value.
// If used on a string, asserts that the string contains value.
func (m *Matcher) Contains(value interface{}) *Matcher {
	// TODO: Implement Matcher.Contains
	return m
}

// Matches asserts that the matched value matches the regex in pattern.
// Note: If the regex cannot be compiled, a fatal error will be thrown.
// If used on a non-string value, a fatal error will be thrown.
func (m *Matcher) Matches(pattern string) *Matcher {
	rv := reflect.ValueOf(m.value)
	if rv.Kind() != reflect.String {
		m.t.Fatalf("expected string, got %s", rv.Kind().String())
	}
	re, err := regexp.Compile(pattern)
	if err != nil {
		m.t.Fatal(err)
	}
	if !re.MatchString(rv.String()) {
		m.t.Errorf("expected %#q to match pattern %#q", rv.String(), pattern)
	}
	return m
}

// KindOf asserts that the matched value of the kind.
func (m *Matcher) KindOf(t reflect.Kind) *Matcher {
	// TODO: Implement Matcher.KindOf
	return m
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package probes

import (
	"encoding/json"
	"net"
	"syscall"
	"time"

	"github.com/vishvananda/netlink"
	"github.com/vishvananda/netlink/nl"

	"github.com/safchain/ethtool"

	"github.com/redhat-cip/skydive/logging"
	"github.com/redhat-cip/skydive/topology/graph"
)

const (
	maxEpollEvents = 32
)

type NetLinkProbe struct {
	Graph             *graph.Graph
	Root              *graph.Node
	nlSocket          *nl.NetlinkSocket
	doneChan          chan struct{}
	indexTointfsQueue map[int64][]*graph.Node
}

func (u *NetLinkProbe) handleIntfIsBridgeMember(intf *graph.Node, link netlink.Link) {
	index := int64(link.Attrs().Index)

	// add children of this interface that haven previously added
	if children, ok := u.indexTointfsQueue[index]; ok {
		for _, child := range children {
			u.Graph.Link(intf, child)
		}
		delete(u.indexTointfsQueue, index)
	}

	// interface being a part of a bridge
	if link.Attrs().MasterIndex != 0 {
		index := int64(link.Attrs().MasterIndex)

		// assuming we have only one parent with this index
		parent := u.Graph.LookupFirstNode(graph.Metadatas{"IfIndex": index, "Type": "bridge"})
		if parent != nil && !u.Graph.AreLinked(parent, intf) {
			u.Graph.Link(parent, intf)
		} else {
			// not yet the bridge so, enqueue for a later add
			u.indexTointfsQueue[index] = append(u.indexTointfsQueue[index], intf)
		}
	}
}

func (u *NetLinkProbe) handleIntfIsVeth(intf *graph.Node, link netlink.Link) {
	if link.Type() != "veth" {
		return
	}

	stats, err := ethtool.Stats(link.Attrs().Name)
	if err != nil {
		logging.GetLogger().Error("Unable get stats from ethtool: %s", err.Error())
		return
	}

	if index, ok := stats["peer_ifindex"]; ok {
		peerResolver := func() bool {
			// re get the interface from the graph since the interface could have been deleted
			if u.Graph.GetNode(intf.ID) == nil {
				return false
			}

			// got more than 1 peer, unable to find the right one, wait for the other to discover
			peer := u.Graph.LookupFirstNode(graph.Metadatas{"IfIndex": int64(index), "Type": "veth"})
			if peer != nil && !u.Graph.AreLinked(peer, intf) {
				u.Graph.NewEdge(graph.GenID(), peer, intf, graph.Metadatas{"Type": "veth"})
				return true
			}
			return false
		}

		if int64(index) > intf.Metadatas()["IfIndex"].(int64) {
			ok := peerResolver()
			if !ok {
				// retry few second later since the right peer can be insert later
				go func() {
					time.Sleep(2 * time.Second)

					u.Graph.Lock()
					defer u.Graph.Unlock()
					peerResolver()
				}()
			}
		}
	}
}

func (u *NetLinkProbe) addGenericLinkToTopology(link netlink.Link, m graph.Metadatas) *graph.Node {
	name := link.Attrs().Name
	index := int64(link.Attrs().Index)

	var intf *graph.Node
	if name != "lo" {
		intf = u.Graph.LookupFirstNode(graph.Metadatas{
			"Name":    name,
			"IfIndex": index,
			"MAC":     link.Attrs().HardwareAddr.String()})
	}

	if intf == nil {
		intf = u.Graph.NewNode(graph.GenID(), m)
	}

	if intf == nil {
		return nil
	}

	if !u.Graph.AreLinked(u.Root, intf) {
		u.Graph.Link(u.Root, intf)
	}

	u.handleIntfIsBridgeMember(intf, link)
	u.handleIntfIsVeth(intf, link)

	return intf
}

func (u *NetLinkProbe) addBridgeLinkToTopology(link netlink.Link, m graph.Metadatas) *graph.Node {
	name := link.Attrs().Name
	index := int64(link.Attrs().Index)

	intf := u.Graph.LookupFirstNode(graph.Metadatas{
		"Name":    name,
		"IfIndex": index})

	if intf == nil {
		intf = u.Graph.NewNode(graph.GenID(), m)
	}

	if !u.Graph.AreLinked(u.Root, intf) {
		u.Graph.Link(u.Root, intf)
	}

	return intf
}

func (u *NetLinkProbe) addOvsLinkToTopology(link netlink.Link, m graph.Metadatas) *graph.Node {
	name := link.Attrs().Name

	intf := u.Graph.LookupFirstNode(graph.Metadatas{"Name": name, "Driver": "openvswitch"})
	if intf == nil {
		intf = u.Graph.NewNode(graph.GenID(), m)
	}

	if !u.Graph.AreLinked(u.Root, intf) {
		u.Graph.Link(u.Root, intf)
	}

	return intf
}

func (u *NetLinkProbe) getLinkIPV4Addr(link netlink.Link) string {
	ipv4 := make([]string, 0)

	addrs, err := netlink.AddrList(link, netlink.FAMILY_V4)
	if err != nil {
		return ""
	}

	for _, addr := range addrs {
		ipv4 = append(ipv4, addr.IPNet.String())
	}

	j, _ := json.Marshal(ipv4)

	return string(j)
}

func (u *NetLinkProbe) addLinkToTopology(link netlink.Link) {
	logging.GetLogger().Debug("Link \"%s(%d)\" added", link.Attrs().Name, link.Attrs().Index)

	u.Graph.Lock()
	defer u.Graph.Unlock()

	driver, _ := ethtool.DriverName(link.Attrs().Name)

	metadatas := graph.Metadatas{
		"Name":    link.Attrs().Name,
		"Type":    link.Type(),
		"IfIndex": int64(link.Attrs().Index),
		"MAC":     link.Attrs().HardwareAddr.String(),
		"MTU":     int64(link.Attrs().MTU),
		"Driver":  driver,
	}

	/*ipv4 := u.getLinkIPV4Addr(link)
	if len(ipv4) > 0 {
		metadatas["IPV4"] = ipv4
	}*/

	if vlan, ok := link.(*netlink.Vlan); ok {
		metadatas["Vlan"] = vlan.VlanId
	}

	if (link.Attrs().Flags & net.FlagUp) > 0 {
		metadatas["State"] = "UP"
	} else {
		metadatas["State"] = "DOWN"
	}

	var intf *graph.Node

	switch driver {
	case "bridge":
		intf = u.addBridgeLinkToTopology(link, metadatas)
	case "openvswitch":
		intf = u.addOvsLinkToTopology(link, metadatas)
	default:
		intf = u.addGenericLinkToTopology(link, metadatas)
	}

	// merge metadatas if the interface returned is not a new one
	if intf != nil {
		m := intf.Metadatas()

		updated := false
		for k, nv := range metadatas {
			if ov, ok := m[k]; ok && nv == ov {
				continue
			}
			m[k] = nv
			updated = true
		}

		if updated {
			u.Graph.SetMetadatas(intf, m)
		}
	}
}

func (u *NetLinkProbe) onLinkAdded(index int) {
	link, err := netlink.LinkByIndex(index)
	if err != nil {
		logging.GetLogger().Error("Failed to find interface %d: %s", index, err.Error())
		return
	}

	u.addLinkToTopology(link)
}

func (u *NetLinkProbe) onLinkDeleted(index int) {
	logging.GetLogger().Debug("Link %d deleted", index)

	u.Graph.Lock()
	defer u.Graph.Unlock()

	var intf *graph.Node

	intfs := u.Graph.LookupNodes(graph.Metadatas{"IfIndex": int64(index)})
	switch l := len(intfs); {
	case l == 1:
		intf = intfs[0]
	case l > 1:
	Loop:
		for _, i := range intfs {
			parents := u.Graph.LookupParentNodes(i, nil)
			for _, parent := range parents {
				if parent.ID == u.Root.ID {
					intf = i
					break Loop
				}
			}
		}
	}

	// case of removing the interface from a bridge
	if intf != nil {
		parents := u.Graph.LookupParentNodes(intf, graph.Metadatas{"Type": "bridge"})
		for _, parent := range parents {
			u.Graph.Unlink(parent, intf)
		}
	}

	// check wheter the interface has been deleted or not
	// we get a delete event when an interace is removed from a bridge
	_, err := netlink.LinkByIndex(index)
	if err != nil && intf != nil {
		if driver, ok := intf.Metadatas()["Driver"]; ok {
			// if openvswitch do not remove let's do the job by ovs piece of code
			if driver == "openvswitch" {
				u.Graph.Unlink(u.Root, intf)
			} else {
				u.Graph.DelNode(intf)
			}
		}
	}

	delete(u.indexTointfsQueue, int64(index))
}

func (u *NetLinkProbe) initialize() {
	links, err := netlink.LinkList()
	if err != nil {
		logging.GetLogger().Error("Unable to list interfaces: %s", err.Error())
		return
	}

	for _, link := range links {
		u.addLinkToTopology(link)
	}
}

func (u *NetLinkProbe) start() {
	s, err := nl.Subscribe(syscall.NETLINK_ROUTE, syscall.RTNLGRP_LINK)
	if err != nil {
		logging.GetLogger().Error("Failed to subscribe to netlink RTNLGRP_LINK messages: %s", err.Error())
		return
	}
	u.nlSocket = s

	fd := u.nlSocket.GetFd()

	err = syscall.SetNonblock(fd, true)
	if err != nil {
		logging.GetLogger().Error("Failed to set the netlink fd as non-blocking: %s", err.Error())
		return
	}

	epfd, e := syscall.EpollCreate1(0)
	if e != nil {
		logging.GetLogger().Error("Failed to set the netlink fd as non-blocking: %s", err.Error())
		return
	}
	defer syscall.Close(epfd)

	u.initialize()

	event := syscall.EpollEvent{Events: syscall.EPOLLIN, Fd: int32(fd)}
	if e = syscall.EpollCtl(epfd, syscall.EPOLL_CTL_ADD, fd, &event); e != nil {
		logging.GetLogger().Error("Failed to set the netlink fd as non-blocking: %s", err.Error())
		return
	}

	events := make([]syscall.EpollEvent, maxEpollEvents)

Loop:
	for {
		n, err := syscall.EpollWait(epfd, events[:], 1000)
		if err != nil {
			logging.GetLogger().Error("Failed to receive from netlink messages: %s", err.Error())
			continue
		}

		if n == 0 {
			select {
			case <-u.doneChan:
				break Loop
			default:
				continue
			}
		}

		msgs, err := s.Receive()
		if err != nil {
			logging.GetLogger().Error("Failed to receive from netlink messages: %s", err.Error())

			time.Sleep(1 * time.Second)
			continue
		}

		for _, msg := range msgs {
			switch msg.Header.Type {
			case syscall.RTM_NEWLINK:
				ifmsg := nl.DeserializeIfInfomsg(msg.Data)
				u.onLinkAdded(int(ifmsg.Index))
			case syscall.RTM_DELLINK:
				ifmsg := nl.DeserializeIfInfomsg(msg.Data)
				u.onLinkDeleted(int(ifmsg.Index))
			}
		}
	}

	u.nlSocket.Close()
}

func (u *NetLinkProbe) Start() {
	go u.start()
}

func (u *NetLinkProbe) Run() {
	u.start()
}

func (u *NetLinkProbe) Stop() {
	u.doneChan <- struct{}{}
}

func NewNetLinkProbe(g *graph.Graph, n *graph.Node) *NetLinkProbe {
	return &NetLinkProbe{
		Graph:             g,
		Root:              n,
		doneChan:          make(chan struct{}),
		indexTointfsQueue: make(map[int64][]*graph.Node),
	}
}
<commit_msg>Force driver to bridge when link type is bridge<commit_after>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package probes

import (
	"encoding/json"
	"net"
	"syscall"
	"time"

	"github.com/vishvananda/netlink"
	"github.com/vishvananda/netlink/nl"

	"github.com/safchain/ethtool"

	"github.com/redhat-cip/skydive/logging"
	"github.com/redhat-cip/skydive/topology/graph"
)

const (
	maxEpollEvents = 32
)

type NetLinkProbe struct {
	Graph             *graph.Graph
	Root              *graph.Node
	nlSocket          *nl.NetlinkSocket
	doneChan          chan struct{}
	indexTointfsQueue map[int64][]*graph.Node
}

func (u *NetLinkProbe) handleIntfIsBridgeMember(intf *graph.Node, link netlink.Link) {
	index := int64(link.Attrs().Index)

	// add children of this interface that haven previously added
	if children, ok := u.indexTointfsQueue[index]; ok {
		for _, child := range children {
			u.Graph.Link(intf, child)
		}
		delete(u.indexTointfsQueue, index)
	}

	// interface being a part of a bridge
	if link.Attrs().MasterIndex != 0 {
		index := int64(link.Attrs().MasterIndex)

		// assuming we have only one parent with this index
		parent := u.Graph.LookupFirstNode(graph.Metadatas{"IfIndex": index, "Type": "bridge"})
		if parent != nil && !u.Graph.AreLinked(parent, intf) {
			u.Graph.Link(parent, intf)
		} else {
			// not yet the bridge so, enqueue for a later add
			u.indexTointfsQueue[index] = append(u.indexTointfsQueue[index], intf)
		}
	}
}

func (u *NetLinkProbe) handleIntfIsVeth(intf *graph.Node, link netlink.Link) {
	if link.Type() != "veth" {
		return
	}

	stats, err := ethtool.Stats(link.Attrs().Name)
	if err != nil {
		logging.GetLogger().Error("Unable get stats from ethtool: %s", err.Error())
		return
	}

	if index, ok := stats["peer_ifindex"]; ok {
		peerResolver := func() bool {
			// re get the interface from the graph since the interface could have been deleted
			if u.Graph.GetNode(intf.ID) == nil {
				return false
			}

			// got more than 1 peer, unable to find the right one, wait for the other to discover
			peer := u.Graph.LookupFirstNode(graph.Metadatas{"IfIndex": int64(index), "Type": "veth"})
			if peer != nil && !u.Graph.AreLinked(peer, intf) {
				u.Graph.NewEdge(graph.GenID(), peer, intf, graph.Metadatas{"Type": "veth"})
				return true
			}
			return false
		}

		if int64(index) > intf.Metadatas()["IfIndex"].(int64) {
			ok := peerResolver()
			if !ok {
				// retry few second later since the right peer can be insert later
				go func() {
					time.Sleep(2 * time.Second)

					u.Graph.Lock()
					defer u.Graph.Unlock()
					peerResolver()
				}()
			}
		}
	}
}

func (u *NetLinkProbe) addGenericLinkToTopology(link netlink.Link, m graph.Metadatas) *graph.Node {
	name := link.Attrs().Name
	index := int64(link.Attrs().Index)

	var intf *graph.Node
	if name != "lo" {
		intf = u.Graph.LookupFirstNode(graph.Metadatas{
			"Name":    name,
			"IfIndex": index,
			"MAC":     link.Attrs().HardwareAddr.String()})
	}

	if intf == nil {
		intf = u.Graph.NewNode(graph.GenID(), m)
	}

	if intf == nil {
		return nil
	}

	if !u.Graph.AreLinked(u.Root, intf) {
		u.Graph.Link(u.Root, intf)
	}

	u.handleIntfIsBridgeMember(intf, link)
	u.handleIntfIsVeth(intf, link)

	return intf
}

func (u *NetLinkProbe) addBridgeLinkToTopology(link netlink.Link, m graph.Metadatas) *graph.Node {
	name := link.Attrs().Name
	index := int64(link.Attrs().Index)

	intf := u.Graph.LookupFirstNode(graph.Metadatas{
		"Name":    name,
		"IfIndex": index})

	if intf == nil {
		intf = u.Graph.NewNode(graph.GenID(), m)
	}

	if !u.Graph.AreLinked(u.Root, intf) {
		u.Graph.Link(u.Root, intf)
	}

	return intf
}

func (u *NetLinkProbe) addOvsLinkToTopology(link netlink.Link, m graph.Metadatas) *graph.Node {
	name := link.Attrs().Name

	intf := u.Graph.LookupFirstNode(graph.Metadatas{"Name": name, "Driver": "openvswitch"})
	if intf == nil {
		intf = u.Graph.NewNode(graph.GenID(), m)
	}

	if !u.Graph.AreLinked(u.Root, intf) {
		u.Graph.Link(u.Root, intf)
	}

	return intf
}

func (u *NetLinkProbe) getLinkIPV4Addr(link netlink.Link) string {
	ipv4 := make([]string, 0)

	addrs, err := netlink.AddrList(link, netlink.FAMILY_V4)
	if err != nil {
		return ""
	}

	for _, addr := range addrs {
		ipv4 = append(ipv4, addr.IPNet.String())
	}

	j, _ := json.Marshal(ipv4)

	return string(j)
}

func (u *NetLinkProbe) addLinkToTopology(link netlink.Link) {
	logging.GetLogger().Debug("Link \"%s(%d)\" added", link.Attrs().Name, link.Attrs().Index)

	u.Graph.Lock()
	defer u.Graph.Unlock()

	driver, _ := ethtool.DriverName(link.Attrs().Name)
	if driver == "" && link.Type() == "bridge" {
		driver = "bridge"
	}

	metadatas := graph.Metadatas{
		"Name":    link.Attrs().Name,
		"Type":    link.Type(),
		"IfIndex": int64(link.Attrs().Index),
		"MAC":     link.Attrs().HardwareAddr.String(),
		"MTU":     int64(link.Attrs().MTU),
		"Driver":  driver,
	}

	/*ipv4 := u.getLinkIPV4Addr(link)
	if len(ipv4) > 0 {
		metadatas["IPV4"] = ipv4
	}*/

	if vlan, ok := link.(*netlink.Vlan); ok {
		metadatas["Vlan"] = vlan.VlanId
	}

	if (link.Attrs().Flags & net.FlagUp) > 0 {
		metadatas["State"] = "UP"
	} else {
		metadatas["State"] = "DOWN"
	}

	var intf *graph.Node

	switch driver {
	case "bridge":
		intf = u.addBridgeLinkToTopology(link, metadatas)
	case "openvswitch":
		intf = u.addOvsLinkToTopology(link, metadatas)
	default:
		intf = u.addGenericLinkToTopology(link, metadatas)
	}

	// merge metadatas if the interface returned is not a new one
	if intf != nil {
		m := intf.Metadatas()

		updated := false
		for k, nv := range metadatas {
			if ov, ok := m[k]; ok && nv == ov {
				continue
			}
			m[k] = nv
			updated = true
		}

		if updated {
			u.Graph.SetMetadatas(intf, m)
		}
	}
}

func (u *NetLinkProbe) onLinkAdded(index int) {
	link, err := netlink.LinkByIndex(index)
	if err != nil {
		logging.GetLogger().Error("Failed to find interface %d: %s", index, err.Error())
		return
	}

	u.addLinkToTopology(link)
}

func (u *NetLinkProbe) onLinkDeleted(index int) {
	logging.GetLogger().Debug("Link %d deleted", index)

	u.Graph.Lock()
	defer u.Graph.Unlock()

	var intf *graph.Node

	intfs := u.Graph.LookupNodes(graph.Metadatas{"IfIndex": int64(index)})
	switch l := len(intfs); {
	case l == 1:
		intf = intfs[0]
	case l > 1:
	Loop:
		for _, i := range intfs {
			parents := u.Graph.LookupParentNodes(i, nil)
			for _, parent := range parents {
				if parent.ID == u.Root.ID {
					intf = i
					break Loop
				}
			}
		}
	}

	// case of removing the interface from a bridge
	if intf != nil {
		parents := u.Graph.LookupParentNodes(intf, graph.Metadatas{"Type": "bridge"})
		for _, parent := range parents {
			u.Graph.Unlink(parent, intf)
		}
	}

	// check wheter the interface has been deleted or not
	// we get a delete event when an interace is removed from a bridge
	_, err := netlink.LinkByIndex(index)
	if err != nil && intf != nil {
		if driver, ok := intf.Metadatas()["Driver"]; ok {
			// if openvswitch do not remove let's do the job by ovs piece of code
			if driver == "openvswitch" {
				u.Graph.Unlink(u.Root, intf)
			} else {
				u.Graph.DelNode(intf)
			}
		}
	}

	delete(u.indexTointfsQueue, int64(index))
}

func (u *NetLinkProbe) initialize() {
	links, err := netlink.LinkList()
	if err != nil {
		logging.GetLogger().Error("Unable to list interfaces: %s", err.Error())
		return
	}

	for _, link := range links {
		u.addLinkToTopology(link)
	}
}

func (u *NetLinkProbe) start() {
	s, err := nl.Subscribe(syscall.NETLINK_ROUTE, syscall.RTNLGRP_LINK)
	if err != nil {
		logging.GetLogger().Error("Failed to subscribe to netlink RTNLGRP_LINK messages: %s", err.Error())
		return
	}
	u.nlSocket = s

	fd := u.nlSocket.GetFd()

	err = syscall.SetNonblock(fd, true)
	if err != nil {
		logging.GetLogger().Error("Failed to set the netlink fd as non-blocking: %s", err.Error())
		return
	}

	epfd, e := syscall.EpollCreate1(0)
	if e != nil {
		logging.GetLogger().Error("Failed to set the netlink fd as non-blocking: %s", err.Error())
		return
	}
	defer syscall.Close(epfd)

	u.initialize()

	event := syscall.EpollEvent{Events: syscall.EPOLLIN, Fd: int32(fd)}
	if e = syscall.EpollCtl(epfd, syscall.EPOLL_CTL_ADD, fd, &event); e != nil {
		logging.GetLogger().Error("Failed to set the netlink fd as non-blocking: %s", err.Error())
		return
	}

	events := make([]syscall.EpollEvent, maxEpollEvents)

Loop:
	for {
		n, err := syscall.EpollWait(epfd, events[:], 1000)
		if err != nil {
			logging.GetLogger().Error("Failed to receive from netlink messages: %s", err.Error())
			continue
		}

		if n == 0 {
			select {
			case <-u.doneChan:
				break Loop
			default:
				continue
			}
		}

		msgs, err := s.Receive()
		if err != nil {
			logging.GetLogger().Error("Failed to receive from netlink messages: %s", err.Error())

			time.Sleep(1 * time.Second)
			continue
		}

		for _, msg := range msgs {
			switch msg.Header.Type {
			case syscall.RTM_NEWLINK:
				ifmsg := nl.DeserializeIfInfomsg(msg.Data)
				u.onLinkAdded(int(ifmsg.Index))
			case syscall.RTM_DELLINK:
				ifmsg := nl.DeserializeIfInfomsg(msg.Data)
				u.onLinkDeleted(int(ifmsg.Index))
			}
		}
	}

	u.nlSocket.Close()
}

func (u *NetLinkProbe) Start() {
	go u.start()
}

func (u *NetLinkProbe) Run() {
	u.start()
}

func (u *NetLinkProbe) Stop() {
	u.doneChan <- struct{}{}
}

func NewNetLinkProbe(g *graph.Graph, n *graph.Node) *NetLinkProbe {
	return &NetLinkProbe{
		Graph:             g,
		Root:              n,
		doneChan:          make(chan struct{}),
		indexTointfsQueue: make(map[int64][]*graph.Node),
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package probes

import (
	"errors"
	"strconv"
	"time"

	"github.com/rackspace/gophercloud"
	"github.com/rackspace/gophercloud/openstack"
	"github.com/rackspace/gophercloud/openstack/networking/v2/extensions/provider"
	"github.com/rackspace/gophercloud/openstack/networking/v2/networks"
	"github.com/rackspace/gophercloud/openstack/networking/v2/ports"
	"github.com/rackspace/gophercloud/pagination"

	"github.com/pmylund/go-cache"

	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/topology/graph"
)

type NeutronMapper struct {
	graph.DefaultGraphListener
	graph           *graph.Graph
	client          *gophercloud.ServiceClient
	cache           *cache.Cache
	nodeUpdaterChan chan graph.Identifier
}

type Attributes struct {
	NetworkID   string
	NetworkName string
	TenantID    string
	VNI         string
}

func (mapper *NeutronMapper) retrievePort(metadata graph.Metadata) (port ports.Port, err error) {
	var opts ports.ListOpts
	var mac string

	/* If we have a MAC address for a device attached to the interface, that is the one that
	 * will be associated with the Neutron port. */
	if attached_mac, ok := metadata["ExtID/attached-mac"]; ok {
		mac = attached_mac.(string)
	} else {
		mac = metadata["MAC"].(string)
	}

	logging.GetLogger().Debugf("Retrieving attributes from Neutron for MAC: %s", mac)

	/* Determine the best way to search for the Neutron port.
	 * We prefer the Neutron port UUID if we have it, but will fall back
	 * to using the MAC address otherwise. */
	if portid, ok := metadata["ExtID/iface-id"]; ok {
		opts.ID = portid.(string)
	} else {
		opts.MACAddress = mac
	}
	pager := ports.List(mapper.client, opts)
	err = pager.EachPage(func(page pagination.Page) (bool, error) {
		portList, err := ports.ExtractPorts(page)
		if err != nil {
			return false, err
		}

		for _, p := range portList {
			if p.MACAddress == mac {
				port = p
				return true, nil
			}
		}

		return true, nil
	})

	if len(port.NetworkID) == 0 {
		return port, errors.New("Unable to find port for MAC address: " + mac)
	}

	return port, err
}

func (mapper *NeutronMapper) retrieveAttributes(metadata graph.Metadata) (*Attributes, error) {
	port, err := mapper.retrievePort(metadata)
	if err != nil {
		return nil, err
	}

	result := networks.Get(mapper.client, port.NetworkID)
	network, err := provider.ExtractGet(result)
	if err != nil {
		return nil, err
	}

	a := &Attributes{
		NetworkID:   port.NetworkID,
		NetworkName: network.Name,
		TenantID:    port.TenantID,
		VNI:         network.SegmentationID,
	}

	return a, nil
}

func (mapper *NeutronMapper) nodeUpdater() {
	logging.GetLogger().Debugf("Starting Neutron updater")
	for nodeID := range mapper.nodeUpdaterChan {
		node := mapper.graph.GetNode(nodeID)
		if node == nil {
			continue
		}

		if _, ok := node.Metadata()["MAC"]; !ok {
			continue
		}

		attrs, err := mapper.retrieveAttributes(node.Metadata())
		if err != nil {
			continue
		}

		mapper.updateNode(node, attrs)
	}
	logging.GetLogger().Debugf("Stopping Neutron updater")
}

func (mapper *NeutronMapper) updateNode(node *graph.Node, attrs *Attributes) {
	mapper.graph.Lock()
	defer mapper.graph.Unlock()

	tr := mapper.graph.StartMetadataTransaction(node)
	defer tr.Commit()

	tr.AddMetadata("Manager", "neutron")

	if attrs.TenantID != "" {
		tr.AddMetadata("Neutron/TenantID", attrs.TenantID)
	}

	if attrs.NetworkID != "" {
		tr.AddMetadata("Neutron/NetworkID", attrs.NetworkID)
	}

	if attrs.NetworkName != "" {
		tr.AddMetadata("Neutron/NetworkName", attrs.NetworkName)
	}

	if segID, err := strconv.Atoi(attrs.VNI); err != nil && segID > 0 {
		tr.AddMetadata("Neutron/VNI", uint64(segID))
	}

	mapper.cache.Set(string(node.ID), attrs, cache.DefaultExpiration)
}

func (mapper *NeutronMapper) EnhanceNode(node *graph.Node) {
	mac, ok := node.Metadata()["MAC"]
	if !ok {
		return
	}

	a, f := mapper.cache.Get(mac.(string))
	if f {
		attrs := a.(Attributes)
		mapper.updateNode(node, &attrs)
		return
	}

	mapper.nodeUpdaterChan <- node.ID
}

func (mapper *NeutronMapper) OnNodeUpdated(n *graph.Node) {
	mapper.EnhanceNode(n)
}

func (mapper *NeutronMapper) OnNodeAdded(n *graph.Node) {
	mapper.EnhanceNode(n)
}

func (mapper *NeutronMapper) Start() {
	go mapper.nodeUpdater()
}

func (mapper *NeutronMapper) Stop() {
	mapper.graph.RemoveEventListener(mapper)
	close(mapper.nodeUpdaterChan)
}

func NewNeutronMapper(g *graph.Graph, authURL string, username string, password string, tenantName string, regionName string) (*NeutronMapper, error) {
	mapper := &NeutronMapper{graph: g}

	opts := gophercloud.AuthOptions{
		IdentityEndpoint: authURL,
		Username:         username,
		Password:         password,
		TenantName:       tenantName,
		AllowReauth:      true,
	}

	provider, err := openstack.AuthenticatedClient(opts)
	if err != nil {
		return nil, err
	}

	/* TODO(safchain) add config param for the Availability */
	client, err := openstack.NewNetworkV2(provider, gophercloud.EndpointOpts{
		Name:         "neutron",
		Region:       regionName,
		Availability: gophercloud.AvailabilityPublic,
	})
	if err != nil {
		return nil, err
	}
	mapper.client = client

	// Create a cache with a default expiration time of 5 minutes, and which
	// purges expired items every 30 seconds
	expire := config.GetConfig().GetInt("cache.expire")
	cleanup := config.GetConfig().GetInt("cache.cleanup")
	mapper.cache = cache.New(time.Duration(expire)*time.Second, time.Duration(cleanup)*time.Second)
	mapper.nodeUpdaterChan = make(chan graph.Identifier, 500)

	g.AddEventListener(mapper)

	return mapper, nil
}

func NewNeutronMapperFromConfig(g *graph.Graph) (*NeutronMapper, error) {
	authURL := config.GetConfig().GetString("openstack.auth_url")
	username := config.GetConfig().GetString("openstack.username")
	password := config.GetConfig().GetString("openstack.password")
	tenantName := config.GetConfig().GetString("openstack.tenant_name")
	regionName := config.GetConfig().GetString("openstack.region_name")

	return NewNeutronMapper(g, authURL, username, password, tenantName, regionName)
}
<commit_msg>Limit neutron API calls<commit_after>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package probes

import (
	"strconv"
	"time"

	"github.com/rackspace/gophercloud"
	"github.com/rackspace/gophercloud/openstack"
	"github.com/rackspace/gophercloud/openstack/networking/v2/extensions/provider"
	"github.com/rackspace/gophercloud/openstack/networking/v2/networks"
	"github.com/rackspace/gophercloud/openstack/networking/v2/ports"
	"github.com/rackspace/gophercloud/pagination"

	"github.com/pmylund/go-cache"

	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/topology/graph"
)

type NeutronMapper struct {
	graph.DefaultGraphListener
	graph           *graph.Graph
	client          *gophercloud.ServiceClient
	cache           *cache.Cache
	nodeUpdaterChan chan graph.Identifier
}

type Attributes struct {
	NetworkID   string
	NetworkName string
	TenantID    string
	VNI         string
}

type NeutronPortNotFound struct {
	MAC string
}

func (e NeutronPortNotFound) Error() string {
	return "Unable to find port for MAC address: " + e.MAC
}

func (mapper *NeutronMapper) retrievePort(metadata graph.Metadata) (port ports.Port, err error) {
	var opts ports.ListOpts
	var mac string

	/* If we have a MAC address for a device attached to the interface, that is the one that
	 * will be associated with the Neutron port. */
	if attached_mac, ok := metadata["ExtID/attached-mac"]; ok {
		mac = attached_mac.(string)
	} else {
		mac = metadata["MAC"].(string)
	}

	logging.GetLogger().Debugf("Retrieving attributes from Neutron for MAC: %s", mac)

	/* Determine the best way to search for the Neutron port.
	 * We prefer the Neutron port UUID if we have it, but will fall back
	 * to using the MAC address otherwise. */
	if portid, ok := metadata["ExtID/iface-id"]; ok {
		opts.ID = portid.(string)
	} else {
		opts.MACAddress = mac
	}
	pager := ports.List(mapper.client, opts)
	err = pager.EachPage(func(page pagination.Page) (bool, error) {
		portList, err := ports.ExtractPorts(page)
		if err != nil {
			return false, err
		}

		for _, p := range portList {
			if p.MACAddress == mac {
				port = p
				return true, nil
			}
		}

		return true, nil
	})

	if len(port.NetworkID) == 0 {
		return port, NeutronPortNotFound{mac}
	}

	return port, err
}

func (mapper *NeutronMapper) retrieveAttributes(metadata graph.Metadata) (*Attributes, error) {
	port, err := mapper.retrievePort(metadata)
	if err != nil {
		return nil, err
	}

	result := networks.Get(mapper.client, port.NetworkID)
	network, err := provider.ExtractGet(result)
	if err != nil {
		return nil, err
	}

	a := &Attributes{
		NetworkID:   port.NetworkID,
		NetworkName: network.Name,
		TenantID:    port.TenantID,
		VNI:         network.SegmentationID,
	}

	return a, nil
}

func (mapper *NeutronMapper) nodeUpdater() {
	logging.GetLogger().Debugf("Starting Neutron updater")
	for nodeID := range mapper.nodeUpdaterChan {
		node := mapper.graph.GetNode(nodeID)
		if node == nil {
			continue
		}

		if _, ok := node.Metadata()["MAC"]; !ok {
			continue
		}

		attrs, err := mapper.retrieveAttributes(node.Metadata())
		if err != nil {
			if nerr, ok := err.(NeutronPortNotFound); ok {
				logging.GetLogger().Debugf("Setting in cache not found MAC " + nerr.MAC)
				mapper.cache.Set(nerr.MAC, nil, cache.DefaultExpiration)
			}
			continue
		}

		mapper.updateNode(node, attrs)
	}
	logging.GetLogger().Debugf("Stopping Neutron updater")
}

func (mapper *NeutronMapper) updateNode(node *graph.Node, attrs *Attributes) {
	mapper.graph.Lock()
	defer mapper.graph.Unlock()

	tr := mapper.graph.StartMetadataTransaction(node)
	defer tr.Commit()

	tr.AddMetadata("Manager", "neutron")

	if attrs.TenantID != "" {
		tr.AddMetadata("Neutron/TenantID", attrs.TenantID)
	}

	if attrs.NetworkID != "" {
		tr.AddMetadata("Neutron/NetworkID", attrs.NetworkID)
	}

	if attrs.NetworkName != "" {
		tr.AddMetadata("Neutron/NetworkName", attrs.NetworkName)
	}

	if segID, err := strconv.Atoi(attrs.VNI); err != nil && segID > 0 {
		tr.AddMetadata("Neutron/VNI", uint64(segID))
	}

	mapper.cache.Set(node.Metadata()["MAC"].(string), attrs, cache.DefaultExpiration)
}

func (mapper *NeutronMapper) EnhanceNode(node *graph.Node) {
	mac, ok := node.Metadata()["MAC"]
	if !ok {
		return
	}

	_, f := mapper.cache.Get(mac.(string))
	if f {
		return
	}

	mapper.nodeUpdaterChan <- node.ID
}

func (mapper *NeutronMapper) OnNodeUpdated(n *graph.Node) {
	mapper.EnhanceNode(n)
}

func (mapper *NeutronMapper) OnNodeAdded(n *graph.Node) {
	mapper.EnhanceNode(n)
}

func (mapper *NeutronMapper) Start() {
	go mapper.nodeUpdater()
}

func (mapper *NeutronMapper) Stop() {
	mapper.graph.RemoveEventListener(mapper)
	close(mapper.nodeUpdaterChan)
}

func NewNeutronMapper(g *graph.Graph, authURL string, username string, password string, tenantName string, regionName string) (*NeutronMapper, error) {
	mapper := &NeutronMapper{graph: g}

	opts := gophercloud.AuthOptions{
		IdentityEndpoint: authURL,
		Username:         username,
		Password:         password,
		TenantName:       tenantName,
		AllowReauth:      true,
	}

	provider, err := openstack.AuthenticatedClient(opts)
	if err != nil {
		return nil, err
	}

	/* TODO(safchain) add config param for the Availability */
	client, err := openstack.NewNetworkV2(provider, gophercloud.EndpointOpts{
		Name:         "neutron",
		Region:       regionName,
		Availability: gophercloud.AvailabilityPublic,
	})
	if err != nil {
		return nil, err
	}
	mapper.client = client

	// Create a cache with a default expiration time of 5 minutes, and which
	// purges expired items every 30 seconds
	expire := config.GetConfig().GetInt("cache.expire")
	cleanup := config.GetConfig().GetInt("cache.cleanup")
	mapper.cache = cache.New(time.Duration(expire)*time.Second, time.Duration(cleanup)*time.Second)
	mapper.nodeUpdaterChan = make(chan graph.Identifier, 500)

	g.AddEventListener(mapper)

	return mapper, nil
}

func NewNeutronMapperFromConfig(g *graph.Graph) (*NeutronMapper, error) {
	authURL := config.GetConfig().GetString("openstack.auth_url")
	username := config.GetConfig().GetString("openstack.username")
	password := config.GetConfig().GetString("openstack.password")
	tenantName := config.GetConfig().GetString("openstack.tenant_name")
	regionName := config.GetConfig().GetString("openstack.region_name")

	return NewNeutronMapper(g, authURL, username, password, tenantName, regionName)
}
<|endoftext|>