<commit_before>package creds

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"sort"
	"strings"

	"github.com/aws/aws-sdk-go/aws/credentials"
)

const (
	consoleTokenURL = "https://signin.%s.com"
)

// Translations defines common mappings for credential variables
var Translations = map[string]map[string]string{
	"envvar": {
		"AWS_ACCESS_KEY_ID":     "AccessKey",
		"AWS_SECRET_ACCESS_KEY": "SecretKey",
		"AWS_SESSION_TOKEN":     "SessionToken",
		"AWS_SECURITY_TOKEN":    "SessionToken",
		"AWS_DEFAULT_REGION":    "Region",
		"AWS_REGION":    "Region",
	},
	"console": {
		"sessionId":    "AccessKey",
		"sessionKey":   "SecretKey",
		"sessionToken": "SessionToken",
	},
}

// Translate converts credentials based on a map of field names
func (c Creds) Translate(dictionary map[string]string) map[string]string {
	logger.InfoMsgf("translating using dictionary: %+v", dictionary)
	old := c.ToMap()
	new := make(map[string]string)
	for k, v := range dictionary {
		new[k] = old[v]
	}
	return new
}

// ToMap returns the credentials as a map of field names to strings
func (c Creds) ToMap() map[string]string {
	logger.InfoMsg("converting credentials to map")
	return map[string]string{
		"AccessKey":    c.AccessKey,
		"SecretKey":    c.SecretKey,
		"SessionToken": c.SessionToken,
		"Region":       c.Region,
	}
}

// ToSdk returns an AWS SDK Credentials object
func (c *Creds) ToSdk() *credentials.Credentials {
	logger.InfoMsg("converting credentials to sdk")
	return credentials.NewStaticCredentials(c.AccessKey, c.SecretKey, c.SessionToken)
}

// ToEnvVars returns environment variables suitable for evaling on the current platform
func (c Creds) ToEnvVars() []string {
	if runtime.GOOS == "windows" {
		return c.ToWindowsEnvVars()
	}
	return c.ToLinuxEnvVars()
}

// ToLinuxEnvVars returns environment variables suitable for eval-ing into the POSIX shell
func (c Creds) ToLinuxEnvVars() []string {
	logger.InfoMsg("converting credentials to linux env vars")
	return c.sprintf("export %s=%s")
}

// ToWindowsEnvVars returns environment variables suitable for eval-ing into Windows Powershell
func (c Creds) ToWindowsEnvVars() []string {
	logger.InfoMsg("converting credentials to windows env vars")
	return c.sprintf("$env:%s = \"%s\"")
}

// ToEnviron returns a golang os.Environ object built from the current env plus these credentials
func (c Creds) ToEnviron() []string {
	logger.InfoMsg("converting credentials to golang-style Environ")
	env := os.Environ()
	var newEnv []string
	for _, item := range env {
		if !strings.HasPrefix(item, "AWS_") {
			newEnv = append(newEnv, item)
		}
	}
	return append(newEnv, c.sprintf("%s=%s")...)
}

func (c Creds) sprintf(fmtStr string) []string {
	envCreds := c.Translate(Translations["envvar"])
	var res []string
	for k, v := range envCreds {
		if v != "" {
			res = append(res, fmt.Sprintf(fmtStr, k, v))
		}
	}
	sort.Strings(res)
	return res
}

type consoleTokenResponse struct {
	SigninToken string
}

func (c Creds) toConsoleToken() (string, error) {
	logger.InfoMsg("generating console token")

	args := []string{"?Action=getSigninToken"}

	consoleCreds := c.Translate(Translations["console"])
	jsonCreds, err := json.Marshal(consoleCreds)
	if err != nil {
		return "", err
	}
	urlCreds := url.QueryEscape(string(jsonCreds))
	paramCreds := fmt.Sprintf("Session=%s", urlCreds)
	args = append(args, paramCreds)

	argString := strings.Join(args, "&")
	namespace, err := c.namespace()
	if err != nil {
		return "", err
	}
	baseURL := fmt.Sprintf(consoleTokenURL, namespace)
	url := strings.Join([]string{baseURL, "/federation", argString}, "")

	logger.InfoMsg("making console token http request")
	resp, err := http.Get(url)

	if err != nil {
		return "", err
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			panic(err)
		}
	}()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	logger.InfoMsg("unmarshalling console token json response")
	tokenObj := consoleTokenResponse{}
	if err := json.Unmarshal(body, &tokenObj); err != nil {
		return "", err
	}

	return tokenObj.SigninToken, nil
}

// ToConsoleURL returns a console URL for the role
func (c Creds) ToConsoleURL() (string, error) {
	return c.ToCustomConsoleURL("")
}

// ToCustomConsoleURL returns a console URL with a custom path
func (c Creds) ToCustomConsoleURL(dest string) (string, error) {
	logger.InfoMsg("generating console url")
	consoleToken, err := c.toConsoleToken()
	if err != nil {
		return "", err
	}
	namespace, err := c.namespace()
	if err != nil {
		return "", err
	}
	baseURL := fmt.Sprintf(consoleTokenURL, namespace)
	var targetURL string
	if c.Region != "" {
		targetURL = fmt.Sprintf("https://%s.console.%s.com/%s", c.Region, namespace, dest)
	} else {
		targetURL = fmt.Sprintf("https://console.%s.com/%s", namespace, dest)
	}
	logger.InfoMsgf("using destination url %s", targetURL)
	urlParts := []string{
		baseURL,
		"/federation",
		"?Action=login",
		"&Issuer=",
		"&Destination=",
		url.QueryEscape(targetURL),
		"&SigninToken=",
		consoleToken,
	}
	urlString := strings.Join(urlParts, "")
	return urlString, nil
}

// ToSignoutURL returns a signout URL for the console
func (c Creds) ToSignoutURL() (string, error) {
	logger.InfoMsg("generating signout url")
	namespace, err := c.namespace()
	if err != nil {
		return "", err
	}
	baseURL := fmt.Sprintf(consoleTokenURL, namespace)
	url := strings.Join([]string{baseURL, "/oauth?Action=logout"}, "")
	return url, nil
}
<commit_msg>Update export.go<commit_after>package creds

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"sort"
	"strings"

	"github.com/aws/aws-sdk-go/aws/credentials"
)

const (
	consoleTokenURL = "https://signin.%s.com"
)

// Translations defines common mappings for credential variables
var Translations = map[string]map[string]string{
	"envvar": {
		"AWS_ACCESS_KEY_ID":     "AccessKey",
		"AWS_SECRET_ACCESS_KEY": "SecretKey",
		"AWS_SESSION_TOKEN":     "SessionToken",
		"AWS_SECURITY_TOKEN":    "SessionToken",
		"AWS_DEFAULT_REGION":    "Region",
		"AWS_REGION":            "Region",
	},
	"console": {
		"sessionId":    "AccessKey",
		"sessionKey":   "SecretKey",
		"sessionToken": "SessionToken",
	},
}

// Translate converts credentials based on a map of field names
func (c Creds) Translate(dictionary map[string]string) map[string]string {
	logger.InfoMsgf("translating using dictionary: %+v", dictionary)
	old := c.ToMap()
	new := make(map[string]string)
	for k, v := range dictionary {
		new[k] = old[v]
	}
	return new
}

// ToMap returns the credentials as a map of field names to strings
func (c Creds) ToMap() map[string]string {
	logger.InfoMsg("converting credentials to map")
	return map[string]string{
		"AccessKey":    c.AccessKey,
		"SecretKey":    c.SecretKey,
		"SessionToken": c.SessionToken,
		"Region":       c.Region,
	}
}

// ToSdk returns an AWS SDK Credentials object
func (c *Creds) ToSdk() *credentials.Credentials {
	logger.InfoMsg("converting credentials to sdk")
	return credentials.NewStaticCredentials(c.AccessKey, c.SecretKey, c.SessionToken)
}

// ToEnvVars returns environment variables suitable for evaling on the current platform
func (c Creds) ToEnvVars() []string {
	if runtime.GOOS == "windows" {
		return c.ToWindowsEnvVars()
	}
	return c.ToLinuxEnvVars()
}

// ToLinuxEnvVars returns environment variables suitable for eval-ing into the POSIX shell
func (c Creds) ToLinuxEnvVars() []string {
	logger.InfoMsg("converting credentials to linux env vars")
	return c.sprintf("export %s=%s")
}

// ToWindowsEnvVars returns environment variables suitable for eval-ing into Windows Powershell
func (c Creds) ToWindowsEnvVars() []string {
	logger.InfoMsg("converting credentials to windows env vars")
	return c.sprintf("$env:%s = \"%s\"")
}

// ToEnviron returns a golang os.Environ object built from the current env plus these credentials
func (c Creds) ToEnviron() []string {
	logger.InfoMsg("converting credentials to golang-style Environ")
	env := os.Environ()
	var newEnv []string
	for _, item := range env {
		if !strings.HasPrefix(item, "AWS_") {
			newEnv = append(newEnv, item)
		}
	}
	return append(newEnv, c.sprintf("%s=%s")...)
}

func (c Creds) sprintf(fmtStr string) []string {
	envCreds := c.Translate(Translations["envvar"])
	var res []string
	for k, v := range envCreds {
		if v != "" {
			res = append(res, fmt.Sprintf(fmtStr, k, v))
		}
	}
	sort.Strings(res)
	return res
}

type consoleTokenResponse struct {
	SigninToken string
}

func (c Creds) toConsoleToken() (string, error) {
	logger.InfoMsg("generating console token")

	args := []string{"?Action=getSigninToken"}

	consoleCreds := c.Translate(Translations["console"])
	jsonCreds, err := json.Marshal(consoleCreds)
	if err != nil {
		return "", err
	}
	urlCreds := url.QueryEscape(string(jsonCreds))
	paramCreds := fmt.Sprintf("Session=%s", urlCreds)
	args = append(args, paramCreds)

	argString := strings.Join(args, "&")
	namespace, err := c.namespace()
	if err != nil {
		return "", err
	}
	baseURL := fmt.Sprintf(consoleTokenURL, namespace)
	url := strings.Join([]string{baseURL, "/federation", argString}, "")

	logger.InfoMsg("making console token http request")
	resp, err := http.Get(url)

	if err != nil {
		return "", err
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			panic(err)
		}
	}()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	logger.InfoMsg("unmarshalling console token json response")
	tokenObj := consoleTokenResponse{}
	if err := json.Unmarshal(body, &tokenObj); err != nil {
		return "", err
	}

	return tokenObj.SigninToken, nil
}

// ToConsoleURL returns a console URL for the role
func (c Creds) ToConsoleURL() (string, error) {
	return c.ToCustomConsoleURL("")
}

// ToCustomConsoleURL returns a console URL with a custom path
func (c Creds) ToCustomConsoleURL(dest string) (string, error) {
	logger.InfoMsg("generating console url")
	consoleToken, err := c.toConsoleToken()
	if err != nil {
		return "", err
	}
	namespace, err := c.namespace()
	if err != nil {
		return "", err
	}
	baseURL := fmt.Sprintf(consoleTokenURL, namespace)
	var targetURL string
	if c.Region != "" {
		targetURL = fmt.Sprintf("https://%s.console.%s.com/%s", c.Region, namespace, dest)
	} else {
		targetURL = fmt.Sprintf("https://console.%s.com/%s", namespace, dest)
	}
	logger.InfoMsgf("using destination url %s", targetURL)
	urlParts := []string{
		baseURL,
		"/federation",
		"?Action=login",
		"&Issuer=",
		"&Destination=",
		url.QueryEscape(targetURL),
		"&SigninToken=",
		consoleToken,
	}
	urlString := strings.Join(urlParts, "")
	return urlString, nil
}

// ToSignoutURL returns a signout URL for the console
func (c Creds) ToSignoutURL() (string, error) {
	logger.InfoMsg("generating signout url")
	namespace, err := c.namespace()
	if err != nil {
		return "", err
	}
	baseURL := fmt.Sprintf(consoleTokenURL, namespace)
	url := strings.Join([]string{baseURL, "/oauth?Action=logout"}, "")
	return url, nil
}
<|endoftext|>


<commit_before>package commands

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/filepathfilter"
	"github.com/git-lfs/git-lfs/git"
	"github.com/git-lfs/git-lfs/git/githistory"
	"github.com/git-lfs/git-lfs/git/odb"
	"github.com/git-lfs/git-lfs/lfs"
	"github.com/git-lfs/git-lfs/tasklog"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/spf13/cobra"
)

func migrateImportCommand(cmd *cobra.Command, args []string) {
	l := tasklog.NewLogger(os.Stderr)
	defer l.Close()

	db, err := getObjectDatabase()
	if err != nil {
		ExitWithError(err)
	}
	defer db.Close()

	if migrateNoRewrite {
		if len(args) == 0 {
			ExitWithError(errors.Errorf("fatal: expected one or more files with --no-rewrite"))
		}

		ref, err := git.CurrentRef()
		if err != nil {
			ExitWithError(errors.Wrap(err, "fatal: unable to find current reference"))
		}

		sha, _ := hex.DecodeString(ref.Sha)
		commit, err := db.Commit(sha)
		if err != nil {
			ExitWithError(errors.Wrap(err, "fatal: unable to load commit"))
		}

		root := commit.TreeID

		filter := git.GetAttributeFilter(cfg.LocalWorkingDir(), cfg.LocalGitDir())
		if len(filter.Include()) == 0 && len(filter.Exclude()) == 0 {
			ExitWithError(errors.Errorf("fatal: no git lfs filters found in .gitattributes"))
		}

		gf := lfs.NewGitFilter(cfg)

		for _, file := range args {
			if filter.Allows(file) {
				root, err = rewriteTree(gf, db, root, file)
				if err != nil {
					ExitWithError(errors.Wrapf(err, "fatal: could not rewrite %q", file))
				}
			} else {
				ExitWithError(errors.Errorf("fatal: file %s did not match any entries in .gitattributes", file))
			}
		}

		name, email := cfg.CurrentCommitter()
		author := fmt.Sprintf("%s <%s>", name, email)

		oid, err := db.WriteCommit(&odb.Commit{
			Author:    author,
			Committer: author,
			ParentIDs: [][]byte{sha},
			Message:   generateMigrateCommitMessage(cmd, strings.Join(args, ",")),
			TreeID:    root,
		})

		if err != nil {
			ExitWithError(errors.Wrap(err, "fatal: unable to write commit"))
		}

		if err := git.UpdateRef(ref, oid, "git lfs migrate import --no-rewrite"); err != nil {
			ExitWithError(errors.Wrap(err, "fatal: unable to update ref"))
		}

		if err := checkoutNonBare(l); err != nil {
			ExitWithError(errors.Wrap(err, "fatal: could not checkout"))
		}

		return
	}

	rewriter := getHistoryRewriter(cmd, db, l)

	tracked := trackedFromFilter(rewriter.Filter())
	exts := tools.NewOrderedSet()
	gitfilter := lfs.NewGitFilter(cfg)

	migrate(args, rewriter, l, &githistory.RewriteOptions{
		Verbose:           migrateVerbose,
		ObjectMapFilePath: objectMapFilePath,
		BlobFn: func(path string, b *odb.Blob) (*odb.Blob, error) {
			if filepath.Base(path) == ".gitattributes" {
				return b, nil
			}

			var buf bytes.Buffer

			if _, err := clean(gitfilter, &buf, b.Contents, path, b.Size); err != nil {
				return nil, err
			}

			if ext := filepath.Ext(path); len(ext) > 0 {
				exts.Add(fmt.Sprintf("*%s filter=lfs diff=lfs merge=lfs -text", ext))
			}

			return &odb.Blob{
				Contents: &buf, Size: int64(buf.Len()),
			}, nil
		},

		TreeCallbackFn: func(path string, t *odb.Tree) (*odb.Tree, error) {
			if path != "/" {
				// Ignore non-root trees.
				return t, nil
			}

			ours := tracked
			if ours.Cardinality() == 0 {
				// If there were no explicitly tracked
				// --include, --exclude filters, assume that the
				// include set is the wildcard filepath
				// extensions of files tracked.
				ours = exts
			}

			theirs, err := trackedFromAttrs(db, t)
			if err != nil {
				return nil, err
			}

			// Create a blob of the attributes that are optionally
			// present in the "t" tree's .gitattributes blob, and
			// union in the patterns that we've tracked.
			//
			// Perform this Union() operation each time we visit a
			// root tree such that if the underlying .gitattributes
			// is present and has a diff between commits in the
			// range of commits to migrate, those changes are
			// preserved.
			blob, err := trackedToBlob(db, theirs.Clone().Union(ours))
			if err != nil {
				return nil, err
			}

			// Finally, return a copy of the tree "t" that has the
			// new .gitattributes file included/replaced.
			return t.Merge(&odb.TreeEntry{
				Name:     ".gitattributes",
				Filemode: 0100644,
				Oid:      blob,
			}), nil
		},

		UpdateRefs: true,
	})

	if err := checkoutNonBare(l); err != nil {
		ExitWithError(errors.Wrap(err, "fatal: could not checkout"))
	}
}

// generateMigrateCommitMessage generates a commit message used with
// --no-rewrite, using --message (if given) or generating one if it isn't.
func generateMigrateCommitMessage(cmd *cobra.Command, patterns string) string {
	if cmd.Flag("message").Changed {
		return migrateCommitMessage
	}
	return fmt.Sprintf("%s: convert to Git LFS", patterns)
}

// checkoutNonBare forces a checkout of the current reference, so long as the
// repository is non-bare.
//
// It returns nil on success, and a non-nil error on failure.
func checkoutNonBare(l *tasklog.Logger) error {
	if bare, _ := git.IsBare(); bare {
		return nil
	}

	t := l.Waiter("migrate: checkout")
	defer t.Complete()

	return git.Checkout("", nil, true)
}

// trackedFromFilter returns an ordered set of strings where each entry is a
// line in the .gitattributes file. It adds/removes the fiter/diff/merge=lfs
// attributes based on patterns included/excldued in the given filter.
func trackedFromFilter(filter *filepathfilter.Filter) *tools.OrderedSet {
	tracked := tools.NewOrderedSet()

	for _, include := range filter.Include() {
		tracked.Add(fmt.Sprintf("%s filter=lfs diff=lfs merge=lfs -text", escapeAttrPattern(include)))
	}

	for _, exclude := range filter.Exclude() {
		tracked.Add(fmt.Sprintf("%s text -filter -merge -diff", escapeAttrPattern(exclude)))
	}

	return tracked
}

var (
	// attrsCache maintains a cache from the hex-encoded SHA1 of a
	// .gitattributes blob to the set of patterns parsed from that blob.
	attrsCache = make(map[string]*tools.OrderedSet)
)

// trackedFromAttrs returns an ordered line-delimited set of the contents of a
// .gitattributes blob in a given tree "t".
//
// It returns an empty set if no attributes file could be found, or an error if
// it could not otherwise be opened.
func trackedFromAttrs(db *odb.ObjectDatabase, t *odb.Tree) (*tools.OrderedSet, error) {
	var oid []byte

	for _, e := range t.Entries {
		if strings.ToLower(e.Name) == ".gitattributes" && e.Type() == odb.BlobObjectType {
			oid = e.Oid
			break
		}
	}

	if oid == nil {
		// TODO(@ttaylorr): make (*tools.OrderedSet)(nil) a valid
		// receiver for non-mutative methods.
		return tools.NewOrderedSet(), nil
	}

	sha1 := hex.EncodeToString(oid)

	if s, ok := attrsCache[sha1]; ok {
		return s, nil
	}

	blob, err := db.Blob(oid)
	if err != nil {
		return nil, err
	}

	attrs := tools.NewOrderedSet()

	scanner := bufio.NewScanner(blob.Contents)
	for scanner.Scan() {
		attrs.Add(scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	attrsCache[sha1] = attrs

	return attrsCache[sha1], nil
}

// trackedToBlob writes and returns the OID of a .gitattributes blob based on
// the patterns given in the ordered set of patterns, "patterns".
func trackedToBlob(db *odb.ObjectDatabase, patterns *tools.OrderedSet) ([]byte, error) {
	var attrs bytes.Buffer

	for pattern := range patterns.Iter() {
		fmt.Fprintf(&attrs, "%s\n", pattern)
	}

	return db.WriteBlob(&odb.Blob{
		Contents: &attrs,
		Size:     int64(attrs.Len()),
	})
}

// rewriteTree replaces the blob at the provided path within the given tree with
// a git lfs pointer. It will recursively rewrite any subtrees along the path to the
// blob.
func rewriteTree(gf *lfs.GitFilter, db *odb.ObjectDatabase, root []byte, path string) ([]byte, error) {
	tree, err := db.Tree(root)
	if err != nil {
		return nil, err
	}

	splits := strings.SplitN(path, "/", 2)

	switch len(splits) {
	case 1:
		// The path points to an entry at the root of this tree, so it must be a blob.
		// Try to replace this blob with a Git LFS pointer.
		index := findEntry(tree, splits[0])
		if index < 0 {
			return nil, errors.Errorf("unable to find entry %s in tree", splits[0])
		}

		blobEntry := tree.Entries[index]
		blob, err := db.Blob(blobEntry.Oid)
		if err != nil {
			return nil, err
		}

		var buf bytes.Buffer

		if _, err := clean(gf, &buf, blob.Contents, blobEntry.Name, blob.Size); err != nil {
			return nil, err
		}

		newOid, err := db.WriteBlob(&odb.Blob{
			Contents: &buf,
			Size:     int64(buf.Len()),
		})

		if err != nil {
			return nil, err
		}

		tree = tree.Merge(&odb.TreeEntry{
			Name:     splits[0],
			Filemode: blobEntry.Filemode,
			Oid:      newOid,
		})
		return db.WriteTree(tree)

	case 2:
		// The path points to an entry in a subtree contained at the root of the tree.
		// Recursively rewrite the subtree.
		head, tail := splits[0], splits[1]

		index := findEntry(tree, head)
		if index < 0 {
			return nil, errors.Errorf("unable to find entry %s in tree", head)
		}

		subtreeEntry := tree.Entries[index]
		if subtreeEntry.Type() != odb.TreeObjectType {
			return nil, errors.Errorf("migrate: expected %s to be a tree, got %s", head, subtreeEntry.Type())
		}

		rewrittenSubtree, err := rewriteTree(gf, db, subtreeEntry.Oid, tail)
		if err != nil {
			return nil, err
		}

		tree = tree.Merge(&odb.TreeEntry{
			Filemode: subtreeEntry.Filemode,
			Name:     subtreeEntry.Name,
			Oid:      rewrittenSubtree,
		})

		return db.WriteTree(tree)

	default:
		return nil, errors.Errorf("error parsing path %s", path)
	}
}

// findEntry searches a tree for the desired entry, and returns the index of that
// entry within the tree's Entries array
func findEntry(t *odb.Tree, name string) int {
	for i, entry := range t.Entries {
		if entry.Name == name {
			return i
		}
	}

	return -1
}
<commit_msg>commands: pre-check files against lfs filters<commit_after>package commands

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/filepathfilter"
	"github.com/git-lfs/git-lfs/git"
	"github.com/git-lfs/git-lfs/git/githistory"
	"github.com/git-lfs/git-lfs/git/odb"
	"github.com/git-lfs/git-lfs/lfs"
	"github.com/git-lfs/git-lfs/tasklog"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/spf13/cobra"
)

func migrateImportCommand(cmd *cobra.Command, args []string) {
	l := tasklog.NewLogger(os.Stderr)
	defer l.Close()

	db, err := getObjectDatabase()
	if err != nil {
		ExitWithError(err)
	}
	defer db.Close()

	if migrateNoRewrite {
		if len(args) == 0 {
			ExitWithError(errors.Errorf("fatal: expected one or more files with --no-rewrite"))
		}

		ref, err := git.CurrentRef()
		if err != nil {
			ExitWithError(errors.Wrap(err, "fatal: unable to find current reference"))
		}

		sha, _ := hex.DecodeString(ref.Sha)
		commit, err := db.Commit(sha)
		if err != nil {
			ExitWithError(errors.Wrap(err, "fatal: unable to load commit"))
		}

		root := commit.TreeID

		filter := git.GetAttributeFilter(cfg.LocalWorkingDir(), cfg.LocalGitDir())
		if len(filter.Include()) == 0 && len(filter.Exclude()) == 0 {
			ExitWithError(errors.Errorf("fatal: no git lfs filters found in .gitattributes"))
		}

		gf := lfs.NewGitFilter(cfg)

		for _, file := range args {
			if !filter.Allows(file) {
				ExitWithError(errors.Errorf("fatal: file %s did not match any entries in .gitattributes", file))
			}
		}

		for _, file := range args {
			root, err = rewriteTree(gf, db, root, file)
			if err != nil {
				ExitWithError(errors.Wrapf(err, "fatal: could not rewrite %q", file))
			}
		}

		name, email := cfg.CurrentCommitter()
		author := fmt.Sprintf("%s <%s>", name, email)

		oid, err := db.WriteCommit(&odb.Commit{
			Author:    author,
			Committer: author,
			ParentIDs: [][]byte{sha},
			Message:   generateMigrateCommitMessage(cmd, strings.Join(args, ",")),
			TreeID:    root,
		})

		if err != nil {
			ExitWithError(errors.Wrap(err, "fatal: unable to write commit"))
		}

		if err := git.UpdateRef(ref, oid, "git lfs migrate import --no-rewrite"); err != nil {
			ExitWithError(errors.Wrap(err, "fatal: unable to update ref"))
		}

		if err := checkoutNonBare(l); err != nil {
			ExitWithError(errors.Wrap(err, "fatal: could not checkout"))
		}

		return
	}

	rewriter := getHistoryRewriter(cmd, db, l)

	tracked := trackedFromFilter(rewriter.Filter())
	exts := tools.NewOrderedSet()
	gitfilter := lfs.NewGitFilter(cfg)

	migrate(args, rewriter, l, &githistory.RewriteOptions{
		Verbose:           migrateVerbose,
		ObjectMapFilePath: objectMapFilePath,
		BlobFn: func(path string, b *odb.Blob) (*odb.Blob, error) {
			if filepath.Base(path) == ".gitattributes" {
				return b, nil
			}

			var buf bytes.Buffer

			if _, err := clean(gitfilter, &buf, b.Contents, path, b.Size); err != nil {
				return nil, err
			}

			if ext := filepath.Ext(path); len(ext) > 0 {
				exts.Add(fmt.Sprintf("*%s filter=lfs diff=lfs merge=lfs -text", ext))
			}

			return &odb.Blob{
				Contents: &buf, Size: int64(buf.Len()),
			}, nil
		},

		TreeCallbackFn: func(path string, t *odb.Tree) (*odb.Tree, error) {
			if path != "/" {
				// Ignore non-root trees.
				return t, nil
			}

			ours := tracked
			if ours.Cardinality() == 0 {
				// If there were no explicitly tracked
				// --include, --exclude filters, assume that the
				// include set is the wildcard filepath
				// extensions of files tracked.
				ours = exts
			}

			theirs, err := trackedFromAttrs(db, t)
			if err != nil {
				return nil, err
			}

			// Create a blob of the attributes that are optionally
			// present in the "t" tree's .gitattributes blob, and
			// union in the patterns that we've tracked.
			//
			// Perform this Union() operation each time we visit a
			// root tree such that if the underlying .gitattributes
			// is present and has a diff between commits in the
			// range of commits to migrate, those changes are
			// preserved.
			blob, err := trackedToBlob(db, theirs.Clone().Union(ours))
			if err != nil {
				return nil, err
			}

			// Finally, return a copy of the tree "t" that has the
			// new .gitattributes file included/replaced.
			return t.Merge(&odb.TreeEntry{
				Name:     ".gitattributes",
				Filemode: 0100644,
				Oid:      blob,
			}), nil
		},

		UpdateRefs: true,
	})

	if err := checkoutNonBare(l); err != nil {
		ExitWithError(errors.Wrap(err, "fatal: could not checkout"))
	}
}

// generateMigrateCommitMessage generates a commit message used with
// --no-rewrite, using --message (if given) or generating one if it isn't.
func generateMigrateCommitMessage(cmd *cobra.Command, patterns string) string {
	if cmd.Flag("message").Changed {
		return migrateCommitMessage
	}
	return fmt.Sprintf("%s: convert to Git LFS", patterns)
}

// checkoutNonBare forces a checkout of the current reference, so long as the
// repository is non-bare.
//
// It returns nil on success, and a non-nil error on failure.
func checkoutNonBare(l *tasklog.Logger) error {
	if bare, _ := git.IsBare(); bare {
		return nil
	}

	t := l.Waiter("migrate: checkout")
	defer t.Complete()

	return git.Checkout("", nil, true)
}

// trackedFromFilter returns an ordered set of strings where each entry is a
// line in the .gitattributes file. It adds/removes the fiter/diff/merge=lfs
// attributes based on patterns included/excldued in the given filter.
func trackedFromFilter(filter *filepathfilter.Filter) *tools.OrderedSet {
	tracked := tools.NewOrderedSet()

	for _, include := range filter.Include() {
		tracked.Add(fmt.Sprintf("%s filter=lfs diff=lfs merge=lfs -text", escapeAttrPattern(include)))
	}

	for _, exclude := range filter.Exclude() {
		tracked.Add(fmt.Sprintf("%s text -filter -merge -diff", escapeAttrPattern(exclude)))
	}

	return tracked
}

var (
	// attrsCache maintains a cache from the hex-encoded SHA1 of a
	// .gitattributes blob to the set of patterns parsed from that blob.
	attrsCache = make(map[string]*tools.OrderedSet)
)

// trackedFromAttrs returns an ordered line-delimited set of the contents of a
// .gitattributes blob in a given tree "t".
//
// It returns an empty set if no attributes file could be found, or an error if
// it could not otherwise be opened.
func trackedFromAttrs(db *odb.ObjectDatabase, t *odb.Tree) (*tools.OrderedSet, error) {
	var oid []byte

	for _, e := range t.Entries {
		if strings.ToLower(e.Name) == ".gitattributes" && e.Type() == odb.BlobObjectType {
			oid = e.Oid
			break
		}
	}

	if oid == nil {
		// TODO(@ttaylorr): make (*tools.OrderedSet)(nil) a valid
		// receiver for non-mutative methods.
		return tools.NewOrderedSet(), nil
	}

	sha1 := hex.EncodeToString(oid)

	if s, ok := attrsCache[sha1]; ok {
		return s, nil
	}

	blob, err := db.Blob(oid)
	if err != nil {
		return nil, err
	}

	attrs := tools.NewOrderedSet()

	scanner := bufio.NewScanner(blob.Contents)
	for scanner.Scan() {
		attrs.Add(scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	attrsCache[sha1] = attrs

	return attrsCache[sha1], nil
}

// trackedToBlob writes and returns the OID of a .gitattributes blob based on
// the patterns given in the ordered set of patterns, "patterns".
func trackedToBlob(db *odb.ObjectDatabase, patterns *tools.OrderedSet) ([]byte, error) {
	var attrs bytes.Buffer

	for pattern := range patterns.Iter() {
		fmt.Fprintf(&attrs, "%s\n", pattern)
	}

	return db.WriteBlob(&odb.Blob{
		Contents: &attrs,
		Size:     int64(attrs.Len()),
	})
}

// rewriteTree replaces the blob at the provided path within the given tree with
// a git lfs pointer. It will recursively rewrite any subtrees along the path to the
// blob.
func rewriteTree(gf *lfs.GitFilter, db *odb.ObjectDatabase, root []byte, path string) ([]byte, error) {
	tree, err := db.Tree(root)
	if err != nil {
		return nil, err
	}

	splits := strings.SplitN(path, "/", 2)

	switch len(splits) {
	case 1:
		// The path points to an entry at the root of this tree, so it must be a blob.
		// Try to replace this blob with a Git LFS pointer.
		index := findEntry(tree, splits[0])
		if index < 0 {
			return nil, errors.Errorf("unable to find entry %s in tree", splits[0])
		}

		blobEntry := tree.Entries[index]
		blob, err := db.Blob(blobEntry.Oid)
		if err != nil {
			return nil, err
		}

		var buf bytes.Buffer

		if _, err := clean(gf, &buf, blob.Contents, blobEntry.Name, blob.Size); err != nil {
			return nil, err
		}

		newOid, err := db.WriteBlob(&odb.Blob{
			Contents: &buf,
			Size:     int64(buf.Len()),
		})

		if err != nil {
			return nil, err
		}

		tree = tree.Merge(&odb.TreeEntry{
			Name:     splits[0],
			Filemode: blobEntry.Filemode,
			Oid:      newOid,
		})
		return db.WriteTree(tree)

	case 2:
		// The path points to an entry in a subtree contained at the root of the tree.
		// Recursively rewrite the subtree.
		head, tail := splits[0], splits[1]

		index := findEntry(tree, head)
		if index < 0 {
			return nil, errors.Errorf("unable to find entry %s in tree", head)
		}

		subtreeEntry := tree.Entries[index]
		if subtreeEntry.Type() != odb.TreeObjectType {
			return nil, errors.Errorf("migrate: expected %s to be a tree, got %s", head, subtreeEntry.Type())
		}

		rewrittenSubtree, err := rewriteTree(gf, db, subtreeEntry.Oid, tail)
		if err != nil {
			return nil, err
		}

		tree = tree.Merge(&odb.TreeEntry{
			Filemode: subtreeEntry.Filemode,
			Name:     subtreeEntry.Name,
			Oid:      rewrittenSubtree,
		})

		return db.WriteTree(tree)

	default:
		return nil, errors.Errorf("error parsing path %s", path)
	}
}

// findEntry searches a tree for the desired entry, and returns the index of that
// entry within the tree's Entries array
func findEntry(t *odb.Tree, name string) int {
	for i, entry := range t.Entries {
		if entry.Name == name {
			return i
		}
	}

	return -1
}
<|endoftext|>


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

package cache

import (
	"crypto"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"

	"go.chromium.org/luci/common/data/text/units"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/system/filesystem"
)

// Cache is a cache of objects.
//
// All implementations must be thread-safe.
type Cache interface {
	io.Closer

	// Keys returns the list of all cached digests in LRU order.
	Keys() isolated.HexDigests

	// Touch updates the LRU position of an item to ensure it is kept in the
	// cache.
	//
	// Returns true if item is in cache.
	Touch(digest isolated.HexDigest) bool

	// Evict removes item from cache if it's there.
	Evict(digest isolated.HexDigest)

	// Add reads data from src and stores it in cache.
	Add(digest isolated.HexDigest, src io.Reader) error

	// AddWithHardlink reads data from src and stores it in cache and hardlink file.
	// This is to avoid file removal by shrink in Add().
	AddWithHardlink(digest isolated.HexDigest, src io.Reader, dest string, perm os.FileMode) error

	// Read returns contents of the cached item.
	Read(digest isolated.HexDigest) (io.ReadCloser, error)

	// Hardlink ensures file at |dest| has the same content as cached |digest|.
	//
	// Note that the behavior when dest already exists is undefined. It will work
	// on all POSIX and may or may not fail on Windows depending on the
	// implementation used. Do not rely on this behavior.
	Hardlink(digest isolated.HexDigest, dest string, perm os.FileMode) error

	// GetAdded returns a list of file size added to cache.
	GetAdded() []int64

	// GetUsed returns a list of file size used from cache.
	GetUsed() []int64
}

// Policies is the policies to use on a cache to limit it's footprint.
//
// It's a cache, not a leak.
type Policies struct {
	// MaxSize trims if the cache gets larger than this value. If 0, the cache is
	// effectively a leak.
	MaxSize units.Size
	// MaxItems is the maximum number of items to keep in the cache. If 0, do not
	// enforce a limit.
	MaxItems int
	// MinFreeSpace trims if disk free space becomes lower than this value.
	// Only makes sense when using disk based cache.
	MinFreeSpace units.Size
}

var ErrInvalidHash = errors.New("invalid hash")

// NewDisk creates a disk based cache.
//
// It may return both a valid Cache and an error if it failed to load the
// previous cache metadata. It is safe to ignore this error. This creates
// cache directory if it doesn't exist.
func NewDisk(policies Policies, path, namespace string) (Cache, error) {
	var err error
	path, err = filepath.Abs(path)
	if err != nil {
		return nil, errors.Annotate(err, "failed to call Abs(%s)", path).Err()
	}
	err = os.MkdirAll(path, 0700)
	if err != nil {
		return nil, errors.Annotate(err, "failed to call MkdirAll(%s)", path).Err()
	}
	d := &disk{
		policies: policies,
		path:     path,
		h:        isolated.GetHash(namespace),
		lru:      makeLRUDict(namespace),
	}
	p := d.statePath()

	err = func() error {
		f, err := os.Open(p)
		if err != nil && os.IsNotExist(err) {
			// The fact that the cache is new is not an error.
			return nil
		}
		if err != nil {
			return err
		}
		defer f.Close()
		return json.NewDecoder(f).Decode(&d.lru)
	}()

	if err != nil {
		// Do not use os.RemoveAll, due to strange 'Access Denied' error on windows
		// in os.MkDir after os.RemoveAll.
		// crbug.com/932396#c123
		files, err := ioutil.ReadDir(path)
		if err != nil {
			return nil, errors.Annotate(err, "failed to call ioutil.ReadDir(%s)", path).Err()
		}

		for _, file := range files {
			p := filepath.Join(path, file.Name())
			if err := os.RemoveAll(p); err != nil {
				return nil, errors.Annotate(err, "failed to call os.RemoveAll(%s)", p).Err()
			}
		}

		d.lru = makeLRUDict(namespace)
	}
	return d, err
}

// Private details.

type disk struct {
	// Immutable.
	policies Policies
	path     string
	h        crypto.Hash

	// Lock protected.
	mu  sync.RWMutex // This protects modification of cached entries under |path| too.
	lru lruDict      // Implements LRU based eviction.

	statsMu sync.Mutex // Protects the stats below
	// TODO(maruel): Add stats about: # removed.
	// TODO(maruel): stateFile
	added []int64
	used  []int64
}

func (d *disk) Close() error {
	d.mu.Lock()
	defer d.mu.Unlock()
	if !d.lru.IsDirty() {
		return nil
	}
	f, err := os.Create(d.statePath())
	if err == nil {
		defer f.Close()
		err = json.NewEncoder(f).Encode(&d.lru)
	}
	return err
}

func (d *disk) Keys() isolated.HexDigests {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.lru.keys()
}

func (d *disk) Touch(digest isolated.HexDigest) bool {
	if !digest.Validate(d.h) {
		return false
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	mtime := time.Now()
	if err := os.Chtimes(d.itemPath(digest), mtime, mtime); err != nil {
		return false
	}
	d.lru.touch(digest)
	return true
}

func (d *disk) Evict(digest isolated.HexDigest) {
	if !digest.Validate(d.h) {
		return
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	d.lru.pop(digest)
	_ = os.Remove(d.itemPath(digest))
}

func (d *disk) Read(digest isolated.HexDigest) (io.ReadCloser, error) {
	if !digest.Validate(d.h) {
		return nil, os.ErrInvalid
	}

	d.mu.Lock()
	f, err := os.Open(d.itemPath(digest))
	if err != nil {
		d.mu.Unlock()
		return nil, err
	}
	d.lru.touch(digest)
	d.mu.Unlock()

	fi, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, errors.Annotate(err, "failed to get stat for %s", digest).Err()
	}

	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	d.used = append(d.used, fi.Size())
	return f, nil
}

func (d *disk) add(digest isolated.HexDigest, src io.Reader, cb func() error) error {
	if !digest.Validate(d.h) {
		return os.ErrInvalid
	}
	tmp, err := ioutil.TempFile(d.path, string(digest)+".*.tmp")
	if err != nil {
		return errors.Annotate(err, "failed to create tempfile for %s", digest).Err()
	}
	// TODO(maruel): Use a LimitedReader flavor that fails when reaching limit.
	h := d.h.New()
	size, err := io.Copy(tmp, io.TeeReader(src, h))
	if err2 := tmp.Close(); err == nil {
		err = err2
	}
	fname := tmp.Name()
	if err != nil {
		_ = os.Remove(fname)
		return err
	}
	if d := isolated.Sum(h); d != digest {
		_ = os.Remove(fname)
		return errors.Annotate(ErrInvalidHash, "invalid hash, got=%s, want=%s", d, digest).Err()
	}
	if units.Size(size) > d.policies.MaxSize {
		_ = os.Remove(fname)
		return errors.Reason("item too large, size=%d, limit=%d", size, d.policies.MaxSize).Err()
	}

	d.mu.Lock()
	defer d.mu.Unlock()

	if err := os.Rename(fname, d.itemPath(digest)); err != nil {
		_ = os.Remove(fname)
		return errors.Annotate(err, "failed to rename %s -> %s", fname, d.itemPath(digest)).Err()
	}

	if cb != nil {
		if err := cb(); err != nil {
			return err
		}
	}

	d.lru.pushFront(digest, units.Size(size))
	if err := d.respectPolicies(); err != nil {
		d.lru.pop(digest)
		return err
	}
	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	d.added = append(d.added, size)
	return nil
}

func (d *disk) Add(digest isolated.HexDigest, src io.Reader) error {
	return d.add(digest, src, nil)
}

func (d *disk) AddWithHardlink(digest isolated.HexDigest, src io.Reader, dest string, perm os.FileMode) error {
	return d.add(digest, src, func() error {
		if err := d.hardlinkUnlocked(digest, dest, perm); err != nil {
			_ = os.Remove(d.itemPath(digest))
			return errors.Annotate(err, "failed to call Hardlink(%s, %s)", digest, dest).Err()
		}
		return nil
	})
}

func (d *disk) Hardlink(digest isolated.HexDigest, dest string, perm os.FileMode) error {
	if runtime.GOOS == "darwin" {
		// Accessing the path, which is being replaced, with os.Link
		// seems to cause flaky 'operation not permitted' failure on
		// macOS (https://crbug.com/1076468). So prevent that by holding
		// read lock here.
		d.mu.RLock()
		defer d.mu.RUnlock()
	}
	return d.hardlinkUnlocked(digest, dest, perm)
}

func (d *disk) hardlinkUnlocked(digest isolated.HexDigest, dest string, perm os.FileMode) error {
	if !digest.Validate(d.h) {
		return os.ErrInvalid
	}
	src := d.itemPath(digest)
	// - Windows, if dest exists, the call fails. In particular, trying to
	//   os.Remove() will fail if the file's ReadOnly bit is set. What's worse is
	//   that the ReadOnly bit is set on the file inode, shared on all hardlinks
	//   to this inode. This means that in the case of a file with the ReadOnly
	//   bit set, it would have to do:
	//   - If dest exists:
	//    - If dest has ReadOnly bit:
	//      - If file has any other inode:
	//        - Remove the ReadOnly bit.
	//        - Remove dest.
	//        - Set the ReadOnly bit on one of the inode found.
	//   - Call os.Link()
	//  In short, nobody ain't got time for that.
	//
	// - On any other (sane) OS, if dest exists, it is silently overwritten.
	if err := os.Link(src, dest); err != nil {
		if _, serr := os.Stat(src); errors.Contains(serr, os.ErrNotExist) {
			// In Windows, os.Link may fail with access denied error even if |src| isn't there.
			// And this is to normalize returned error in such case.
			// https://crbug.com/1098265
			err = errors.Annotate(serr, "%s doesn't exist and os.Link failed: %v", src, err).Err()
		}
		debugInfo := fmt.Sprintf("Stats:\n*  src: %s\n*  dest: %s\n*  destDir: %s\nUID=%d GID=%d", statsStr(src), statsStr(dest), statsStr(filepath.Dir(dest)), os.Getuid(), os.Getgid())
		return errors.Annotate(err, "failed to call os.Link(%s, %s)\n%s", src, dest, debugInfo).Err()
	}

	if err := os.Chmod(dest, perm); err != nil {
		return errors.Annotate(err, "failed to call os.Chmod(%s, %#o)", dest, perm).Err()
	}

	fi, err := os.Stat(dest)
	if err != nil {
		return errors.Annotate(err, "failed to call os.Stat(%s)", dest).Err()
	}
	size := fi.Size()
	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	// If this succeeds directly, it means the file is already cached on the
	// disk, so we put it into LRU.
	d.used = append(d.used, size)

	return nil
}

func (d *disk) GetAdded() []int64 {
	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	return append([]int64{}, d.added...)
}

func (d *disk) GetUsed() []int64 {
	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	return append([]int64{}, d.used...)
}

func (d *disk) itemPath(digest isolated.HexDigest) string {
	return filepath.Join(d.path, string(digest))
}

func (d *disk) statePath() string {
	return filepath.Join(d.path, "state.json")
}

func (d *disk) respectPolicies() error {
	minFreeSpaceWanted := uint64(d.policies.MinFreeSpace)
	for {
		freeSpace, err := filesystem.GetFreeSpace(d.path)
		if err != nil {
			return errors.Annotate(err, "couldn't estimate the free space at %s", d.path).Err()
		}
		if d.lru.length() <= d.policies.MaxItems && d.lru.sum <= d.policies.MaxSize && freeSpace >= minFreeSpaceWanted {
			break
		}
		if d.lru.length() == 0 {
			return errors.Reason("no more space to free: current free space=%d policies.MinFreeSpace=%d", freeSpace, minFreeSpaceWanted).Err()
		}
		k, _ := d.lru.popOldest()
		_ = os.Remove(d.itemPath(k))
	}
	return nil
}

func statsStr(path string) string {
	fi, err := os.Stat(path)
	return fmt.Sprintf("path=%s FileInfo=%+v err=%v", path, fi, err)
}
<commit_msg>cache: use exclusive lock in hardlink on macOS<commit_after>// Copyright 2015 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cache

import (
	"crypto"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"

	"go.chromium.org/luci/common/data/text/units"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/system/filesystem"
)

// Cache is a cache of objects.
//
// All implementations must be thread-safe.
type Cache interface {
	io.Closer

	// Keys returns the list of all cached digests in LRU order.
	Keys() isolated.HexDigests

	// Touch updates the LRU position of an item to ensure it is kept in the
	// cache.
	//
	// Returns true if item is in cache.
	Touch(digest isolated.HexDigest) bool

	// Evict removes item from cache if it's there.
	Evict(digest isolated.HexDigest)

	// Add reads data from src and stores it in cache.
	Add(digest isolated.HexDigest, src io.Reader) error

	// AddWithHardlink reads data from src and stores it in cache and hardlink file.
	// This is to avoid file removal by shrink in Add().
	AddWithHardlink(digest isolated.HexDigest, src io.Reader, dest string, perm os.FileMode) error

	// Read returns contents of the cached item.
	Read(digest isolated.HexDigest) (io.ReadCloser, error)

	// Hardlink ensures file at |dest| has the same content as cached |digest|.
	//
	// Note that the behavior when dest already exists is undefined. It will work
	// on all POSIX and may or may not fail on Windows depending on the
	// implementation used. Do not rely on this behavior.
	Hardlink(digest isolated.HexDigest, dest string, perm os.FileMode) error

	// GetAdded returns a list of file size added to cache.
	GetAdded() []int64

	// GetUsed returns a list of file size used from cache.
	GetUsed() []int64
}

// Policies is the policies to use on a cache to limit it's footprint.
//
// It's a cache, not a leak.
type Policies struct {
	// MaxSize trims if the cache gets larger than this value. If 0, the cache is
	// effectively a leak.
	MaxSize units.Size
	// MaxItems is the maximum number of items to keep in the cache. If 0, do not
	// enforce a limit.
	MaxItems int
	// MinFreeSpace trims if disk free space becomes lower than this value.
	// Only makes sense when using disk based cache.
	MinFreeSpace units.Size
}

var ErrInvalidHash = errors.New("invalid hash")

// NewDisk creates a disk based cache.
//
// It may return both a valid Cache and an error if it failed to load the
// previous cache metadata. It is safe to ignore this error. This creates
// cache directory if it doesn't exist.
func NewDisk(policies Policies, path, namespace string) (Cache, error) {
	var err error
	path, err = filepath.Abs(path)
	if err != nil {
		return nil, errors.Annotate(err, "failed to call Abs(%s)", path).Err()
	}
	err = os.MkdirAll(path, 0700)
	if err != nil {
		return nil, errors.Annotate(err, "failed to call MkdirAll(%s)", path).Err()
	}
	d := &disk{
		policies: policies,
		path:     path,
		h:        isolated.GetHash(namespace),
		lru:      makeLRUDict(namespace),
	}
	p := d.statePath()

	err = func() error {
		f, err := os.Open(p)
		if err != nil && os.IsNotExist(err) {
			// The fact that the cache is new is not an error.
			return nil
		}
		if err != nil {
			return err
		}
		defer f.Close()
		return json.NewDecoder(f).Decode(&d.lru)
	}()

	if err != nil {
		// Do not use os.RemoveAll, due to strange 'Access Denied' error on windows
		// in os.MkDir after os.RemoveAll.
		// crbug.com/932396#c123
		files, err := ioutil.ReadDir(path)
		if err != nil {
			return nil, errors.Annotate(err, "failed to call ioutil.ReadDir(%s)", path).Err()
		}

		for _, file := range files {
			p := filepath.Join(path, file.Name())
			if err := os.RemoveAll(p); err != nil {
				return nil, errors.Annotate(err, "failed to call os.RemoveAll(%s)", p).Err()
			}
		}

		d.lru = makeLRUDict(namespace)
	}
	return d, err
}

// Private details.

type disk struct {
	// Immutable.
	policies Policies
	path     string
	h        crypto.Hash

	// Lock protected.
	mu  sync.Mutex // This protects modification of cached entries under |path| too.
	lru lruDict    // Implements LRU based eviction.

	statsMu sync.Mutex // Protects the stats below
	// TODO(maruel): Add stats about: # removed.
	// TODO(maruel): stateFile
	added []int64
	used  []int64
}

func (d *disk) Close() error {
	d.mu.Lock()
	defer d.mu.Unlock()
	if !d.lru.IsDirty() {
		return nil
	}
	f, err := os.Create(d.statePath())
	if err == nil {
		defer f.Close()
		err = json.NewEncoder(f).Encode(&d.lru)
	}
	return err
}

func (d *disk) Keys() isolated.HexDigests {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.lru.keys()
}

func (d *disk) Touch(digest isolated.HexDigest) bool {
	if !digest.Validate(d.h) {
		return false
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	mtime := time.Now()
	if err := os.Chtimes(d.itemPath(digest), mtime, mtime); err != nil {
		return false
	}
	d.lru.touch(digest)
	return true
}

func (d *disk) Evict(digest isolated.HexDigest) {
	if !digest.Validate(d.h) {
		return
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	d.lru.pop(digest)
	_ = os.Remove(d.itemPath(digest))
}

func (d *disk) Read(digest isolated.HexDigest) (io.ReadCloser, error) {
	if !digest.Validate(d.h) {
		return nil, os.ErrInvalid
	}

	d.mu.Lock()
	f, err := os.Open(d.itemPath(digest))
	if err != nil {
		d.mu.Unlock()
		return nil, err
	}
	d.lru.touch(digest)
	d.mu.Unlock()

	fi, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, errors.Annotate(err, "failed to get stat for %s", digest).Err()
	}

	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	d.used = append(d.used, fi.Size())
	return f, nil
}

func (d *disk) add(digest isolated.HexDigest, src io.Reader, cb func() error) error {
	if !digest.Validate(d.h) {
		return os.ErrInvalid
	}
	tmp, err := ioutil.TempFile(d.path, string(digest)+".*.tmp")
	if err != nil {
		return errors.Annotate(err, "failed to create tempfile for %s", digest).Err()
	}
	// TODO(maruel): Use a LimitedReader flavor that fails when reaching limit.
	h := d.h.New()
	size, err := io.Copy(tmp, io.TeeReader(src, h))
	if err2 := tmp.Close(); err == nil {
		err = err2
	}
	fname := tmp.Name()
	if err != nil {
		_ = os.Remove(fname)
		return err
	}
	if d := isolated.Sum(h); d != digest {
		_ = os.Remove(fname)
		return errors.Annotate(ErrInvalidHash, "invalid hash, got=%s, want=%s", d, digest).Err()
	}
	if units.Size(size) > d.policies.MaxSize {
		_ = os.Remove(fname)
		return errors.Reason("item too large, size=%d, limit=%d", size, d.policies.MaxSize).Err()
	}

	d.mu.Lock()
	defer d.mu.Unlock()

	if err := os.Rename(fname, d.itemPath(digest)); err != nil {
		_ = os.Remove(fname)
		return errors.Annotate(err, "failed to rename %s -> %s", fname, d.itemPath(digest)).Err()
	}

	if cb != nil {
		if err := cb(); err != nil {
			return err
		}
	}

	d.lru.pushFront(digest, units.Size(size))
	if err := d.respectPolicies(); err != nil {
		d.lru.pop(digest)
		return err
	}
	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	d.added = append(d.added, size)
	return nil
}

func (d *disk) Add(digest isolated.HexDigest, src io.Reader) error {
	return d.add(digest, src, nil)
}

func (d *disk) AddWithHardlink(digest isolated.HexDigest, src io.Reader, dest string, perm os.FileMode) error {
	return d.add(digest, src, func() error {
		if err := d.hardlinkUnlocked(digest, dest, perm); err != nil {
			_ = os.Remove(d.itemPath(digest))
			return errors.Annotate(err, "failed to call Hardlink(%s, %s)", digest, dest).Err()
		}
		return nil
	})
}

func (d *disk) Hardlink(digest isolated.HexDigest, dest string, perm os.FileMode) error {
	if runtime.GOOS == "darwin" {
		// Accessing the path, which is being replaced, with os.Link
		// seems to cause flaky 'operation not permitted' failure on
		// macOS (https://crbug.com/1076468). So prevent that by holding
		// lock here.
		d.mu.Lock()
		defer d.mu.Unlock()
	}
	return d.hardlinkUnlocked(digest, dest, perm)
}

func (d *disk) hardlinkUnlocked(digest isolated.HexDigest, dest string, perm os.FileMode) error {
	if !digest.Validate(d.h) {
		return os.ErrInvalid
	}
	src := d.itemPath(digest)
	// - Windows, if dest exists, the call fails. In particular, trying to
	//   os.Remove() will fail if the file's ReadOnly bit is set. What's worse is
	//   that the ReadOnly bit is set on the file inode, shared on all hardlinks
	//   to this inode. This means that in the case of a file with the ReadOnly
	//   bit set, it would have to do:
	//   - If dest exists:
	//    - If dest has ReadOnly bit:
	//      - If file has any other inode:
	//        - Remove the ReadOnly bit.
	//        - Remove dest.
	//        - Set the ReadOnly bit on one of the inode found.
	//   - Call os.Link()
	//  In short, nobody ain't got time for that.
	//
	// - On any other (sane) OS, if dest exists, it is silently overwritten.
	if err := os.Link(src, dest); err != nil {
		if _, serr := os.Stat(src); errors.Contains(serr, os.ErrNotExist) {
			// In Windows, os.Link may fail with access denied error even if |src| isn't there.
			// And this is to normalize returned error in such case.
			// https://crbug.com/1098265
			err = errors.Annotate(serr, "%s doesn't exist and os.Link failed: %v", src, err).Err()
		}
		debugInfo := fmt.Sprintf("Stats:\n*  src: %s\n*  dest: %s\n*  destDir: %s\nUID=%d GID=%d", statsStr(src), statsStr(dest), statsStr(filepath.Dir(dest)), os.Getuid(), os.Getgid())
		return errors.Annotate(err, "failed to call os.Link(%s, %s)\n%s", src, dest, debugInfo).Err()
	}

	if err := os.Chmod(dest, perm); err != nil {
		return errors.Annotate(err, "failed to call os.Chmod(%s, %#o)", dest, perm).Err()
	}

	fi, err := os.Stat(dest)
	if err != nil {
		return errors.Annotate(err, "failed to call os.Stat(%s)", dest).Err()
	}
	size := fi.Size()
	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	// If this succeeds directly, it means the file is already cached on the
	// disk, so we put it into LRU.
	d.used = append(d.used, size)

	return nil
}

func (d *disk) GetAdded() []int64 {
	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	return append([]int64{}, d.added...)
}

func (d *disk) GetUsed() []int64 {
	d.statsMu.Lock()
	defer d.statsMu.Unlock()
	return append([]int64{}, d.used...)
}

func (d *disk) itemPath(digest isolated.HexDigest) string {
	return filepath.Join(d.path, string(digest))
}

func (d *disk) statePath() string {
	return filepath.Join(d.path, "state.json")
}

func (d *disk) respectPolicies() error {
	minFreeSpaceWanted := uint64(d.policies.MinFreeSpace)
	for {
		freeSpace, err := filesystem.GetFreeSpace(d.path)
		if err != nil {
			return errors.Annotate(err, "couldn't estimate the free space at %s", d.path).Err()
		}
		if d.lru.length() <= d.policies.MaxItems && d.lru.sum <= d.policies.MaxSize && freeSpace >= minFreeSpaceWanted {
			break
		}
		if d.lru.length() == 0 {
			return errors.Reason("no more space to free: current free space=%d policies.MinFreeSpace=%d", freeSpace, minFreeSpaceWanted).Err()
		}
		k, _ := d.lru.popOldest()
		_ = os.Remove(d.itemPath(k))
	}
	return nil
}

func statsStr(path string) string {
	fi, err := os.Stat(path)
	return fmt.Sprintf("path=%s FileInfo=%+v err=%v", path, fi, err)
}
<|endoftext|>


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

package agent

import (
	"encoding/json"
	"net"
)

const (
	// Name of the option (see Options field in
	// NetIf below) that specifies Kubernetes
	// namespace isolation value (on/off).
	namespaceIsolationOption = "namespace_isolation"
)

// NetworkRequest specifies messages sent to the
// agent containing information on how to configure network
// on its host.
type NetworkRequest struct {
	NetIf NetIf `json:"net_if,omitempty"`
	// TODO we should not need this tag
	Options map[string]string `json:"options,omitempty"`
}

// NetIf is a structure that represents
// network interface and its IP configuration
// together with basic methods operating on this structure.
type NetIf struct {
	Name string `form:"interface_name" json:"interface_name"`
	Mac  string `form:"mac_address,omitempty" json:"interface_name,omitempty"`
	IP   net.IP `form:"ip_address,omitempty" json:"ip_address,omitempty"`
}

// SetIP parses and sets the IP address of the interface.
func (netif *NetIf) SetIP(ip string) error {
	netif.IP = net.ParseIP(ip)
	if netif.IP == nil {
		return failedToParseNetif()
	}
	return nil
}

// UnmarshalJSON results in having NetIf implement Unmarshaler
// interface from encoding/json. This is needed because we use
// a type like net.IP here, not a simple type, and so a call to
// net.ParseIP is required to unmarshal this properly.
func (netif *NetIf) UnmarshalJSON(data []byte) error {
	m := make(map[string]string)
	json.Unmarshal(data, &m)

	netif.IP = net.ParseIP(m["ip_address"])
	if netif.IP == nil {
		return failedToParseNetif()
	}

	netif.Name = m["interface_name"]
	netif.Mac = m["mac_address"]
	return nil
}
<commit_msg>The 'form' part can't actually handle options like omitempty<commit_after>// Copyright (c) 2016 Pani Networks
// All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package agent

import (
	"encoding/json"
	"net"
)

const (
	// Name of the option (see Options field in
	// NetIf below) that specifies Kubernetes
	// namespace isolation value (on/off).
	namespaceIsolationOption = "namespace_isolation"
)

// NetworkRequest specifies messages sent to the
// agent containing information on how to configure network
// on its host.
type NetworkRequest struct {
	NetIf NetIf `json:"net_if,omitempty"`
	// TODO we should not need this tag
	Options map[string]string `json:"options,omitempty"`
}

// NetIf is a structure that represents
// network interface and its IP configuration
// together with basic methods operating on this structure.
type NetIf struct {
	Name string `form:"interface_name" json:"interface_name"`
	Mac  string `form:"mac_address" json:"interface_name,omitempty"`
	IP   net.IP `form:"ip_address" json:"ip_address,omitempty"`
}

// SetIP parses and sets the IP address of the interface.
func (netif *NetIf) SetIP(ip string) error {
	netif.IP = net.ParseIP(ip)
	if netif.IP == nil {
		return failedToParseNetif()
	}
	return nil
}

// UnmarshalJSON results in having NetIf implement Unmarshaler
// interface from encoding/json. This is needed because we use
// a type like net.IP here, not a simple type, and so a call to
// net.ParseIP is required to unmarshal this properly.
func (netif *NetIf) UnmarshalJSON(data []byte) error {
	m := make(map[string]string)
	json.Unmarshal(data, &m)

	netif.IP = net.ParseIP(m["ip_address"])
	if netif.IP == nil {
		return failedToParseNetif()
	}

	netif.Name = m["interface_name"]
	netif.Mac = m["mac_address"]
	return nil
}
<|endoftext|>


<commit_before>package ctrie

import (
	"strconv"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestInsertAndLookup(t *testing.T) {
	assert := assert.New(t)
	ctrie := New()

	_, ok := ctrie.Lookup([]byte("foo"))
	assert.False(ok)

	ctrie.Insert([]byte("foo"), "bar")
	val, ok := ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)

	ctrie.Insert([]byte("fooooo"), "baz")
	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)
	val, ok = ctrie.Lookup([]byte("fooooo"))
	assert.True(ok)
	assert.Equal("baz", val)

	for i := 0; i < 100; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), "blah")
	}
	for i := 0; i < 100; i++ {
		val, ok = ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal("blah", val)
	}

	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)
	ctrie.Insert([]byte("foo"), "qux")
	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("qux", val)
}
<commit_msg>Add ctrie benchmarks<commit_after>package ctrie

import (
	"strconv"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestInsertAndLookup(t *testing.T) {
	assert := assert.New(t)
	ctrie := New()

	_, ok := ctrie.Lookup([]byte("foo"))
	assert.False(ok)

	ctrie.Insert([]byte("foo"), "bar")
	val, ok := ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)

	ctrie.Insert([]byte("fooooo"), "baz")
	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)
	val, ok = ctrie.Lookup([]byte("fooooo"))
	assert.True(ok)
	assert.Equal("baz", val)

	for i := 0; i < 100; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), "blah")
	}
	for i := 0; i < 100; i++ {
		val, ok = ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal("blah", val)
	}

	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)
	ctrie.Insert([]byte("foo"), "qux")
	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("qux", val)
}

func BenchmarkInsert(b *testing.B) {
	ctrie := New()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ctrie.Insert([]byte("foo"), 0)
	}
}

func BenchmarkLookup(b *testing.B) {
	numItems := 1000
	ctrie := New()
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	key := []byte(strconv.Itoa(numItems / 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.Lookup(key)
	}
}
<|endoftext|>


<commit_before>// run

// Copyright 2013 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

func main() {
	ok := true
	for _, tt := range tests {
		func() {
			defer func() {
				if err := recover(); err == nil {
					println(tt.name, "did not panic")
					ok = false
				}
			}()
			tt.fn()
		}()
	}
	if !ok {
		println("BUG")
	}
}

var intp *int
var slicep *[]byte
var a10p *[10]int
var a10Mp *[1<<20]int
var structp *Struct
var bigstructp *BigStruct
var i int
var m *M
var m1 *M1
var m2 *M2

func use(interface{}) {
}

var tests = []struct{
	name string
	fn func()
}{
	// Edit .+1,/^}/s/^[^	].+/	{"&", func() { println(&) }},\n	{"\&&", func() { println(\&&) }},/g
	{"*intp", func() { println(*intp) }},
	{"&*intp", func() { println(&*intp) }},
	{"*slicep", func() { println(*slicep) }},
	{"&*slicep", func() { println(&*slicep) }},
	{"(*slicep)[0]", func() { println((*slicep)[0]) }},
	{"&(*slicep)[0]", func() { println(&(*slicep)[0]) }},
	{"(*slicep)[i]", func() { println((*slicep)[i]) }},
	{"&(*slicep)[i]", func() { println(&(*slicep)[i]) }},
	{"*a10p", func() { use(*a10p) }},
	{"&*a10p", func() { println(&*a10p) }},
	{"a10p[0]", func() { println(a10p[0]) }},
	{"&a10p[0]", func() { println(&a10p[0]) }},
	{"a10p[i]", func() { println(a10p[i]) }},
	{"&a10p[i]", func() { println(&a10p[i]) }},
	{"*structp", func() { use(*structp) }},
	{"&*structp", func() { println(&*structp) }},
	{"structp.i", func() { println(structp.i) }},
	{"&structp.i", func() { println(&structp.i) }},
	{"structp.j", func() { println(structp.j) }},
	{"&structp.j", func() { println(&structp.j) }},
	{"structp.k", func() { println(structp.k) }},
	{"&structp.k", func() { println(&structp.k) }},
	{"structp.x[0]", func() { println(structp.x[0]) }},
	{"&structp.x[0]", func() { println(&structp.x[0]) }},
	{"structp.x[i]", func() { println(structp.x[i]) }},
	{"&structp.x[i]", func() { println(&structp.x[i]) }},
	{"structp.x[9]", func() { println(structp.x[9]) }},
	{"&structp.x[9]", func() { println(&structp.x[9]) }},
	{"structp.l", func() { println(structp.l) }},
	{"&structp.l", func() { println(&structp.l) }},
	{"*bigstructp", func() { use(*bigstructp) }},
	{"&*bigstructp", func() { println(&*bigstructp) }},
	{"bigstructp.i", func() { println(bigstructp.i) }},
	{"&bigstructp.i", func() { println(&bigstructp.i) }},
	{"bigstructp.j", func() { println(bigstructp.j) }},
	{"&bigstructp.j", func() { println(&bigstructp.j) }},
	{"bigstructp.k", func() { println(bigstructp.k) }},
	{"&bigstructp.k", func() { println(&bigstructp.k) }},
	{"bigstructp.x[0]", func() { println(bigstructp.x[0]) }},
	{"&bigstructp.x[0]", func() { println(&bigstructp.x[0]) }},
	{"bigstructp.x[i]", func() { println(bigstructp.x[i]) }},
	{"&bigstructp.x[i]", func() { println(&bigstructp.x[i]) }},
	{"bigstructp.x[9]", func() { println(bigstructp.x[9]) }},
	{"&bigstructp.x[9]", func() { println(&bigstructp.x[9]) }},
	{"bigstructp.x[100<<20]", func() { println(bigstructp.x[100<<20]) }},
	{"&bigstructp.x[100<<20]", func() { println(&bigstructp.x[100<<20]) }},
	{"bigstructp.l", func() { println(bigstructp.l) }},
	{"&bigstructp.l", func() { println(&bigstructp.l) }},
	{"m1.F()", func() { println(m1.F()) }},
	{"m1.M.F()", func() { println(m1.M.F()) }},
	{"m2.F()", func() { println(m2.F()) }},
	{"m2.M.F()", func() { println(m2.M.F()) }},
}

type Struct struct {
	i int
	j float64
	k string
	x [10]int
	l []byte
}

type BigStruct struct {
	i int
	j float64
	k string
	x [128<<20]byte
	l []byte
}

type M struct {
}

func (m *M) F() int {return 0}

type M1 struct {
	M
}

type M2 struct {
	x int
	M
}
<commit_msg>test: exit non-zero on error from nilptr2.go.<commit_after>// run

// Copyright 2013 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import "os"

func main() {
	ok := true
	for _, tt := range tests {
		func() {
			defer func() {
				if err := recover(); err == nil {
					println(tt.name, "did not panic")
					ok = false
				}
			}()
			tt.fn()
		}()
	}
	if !ok {
		println("BUG")
		os.Exit(1)
	}
}

var intp *int
var slicep *[]byte
var a10p *[10]int
var a10Mp *[1<<20]int
var structp *Struct
var bigstructp *BigStruct
var i int
var m *M
var m1 *M1
var m2 *M2

func use(interface{}) {
}

var tests = []struct{
	name string
	fn func()
}{
	// Edit .+1,/^}/s/^[^	].+/	{"&", func() { println(&) }},\n	{"\&&", func() { println(\&&) }},/g
	{"*intp", func() { println(*intp) }},
	{"&*intp", func() { println(&*intp) }},
	{"*slicep", func() { println(*slicep) }},
	{"&*slicep", func() { println(&*slicep) }},
	{"(*slicep)[0]", func() { println((*slicep)[0]) }},
	{"&(*slicep)[0]", func() { println(&(*slicep)[0]) }},
	{"(*slicep)[i]", func() { println((*slicep)[i]) }},
	{"&(*slicep)[i]", func() { println(&(*slicep)[i]) }},
	{"*a10p", func() { use(*a10p) }},
	{"&*a10p", func() { println(&*a10p) }},
	{"a10p[0]", func() { println(a10p[0]) }},
	{"&a10p[0]", func() { println(&a10p[0]) }},
	{"a10p[i]", func() { println(a10p[i]) }},
	{"&a10p[i]", func() { println(&a10p[i]) }},
	{"*structp", func() { use(*structp) }},
	{"&*structp", func() { println(&*structp) }},
	{"structp.i", func() { println(structp.i) }},
	{"&structp.i", func() { println(&structp.i) }},
	{"structp.j", func() { println(structp.j) }},
	{"&structp.j", func() { println(&structp.j) }},
	{"structp.k", func() { println(structp.k) }},
	{"&structp.k", func() { println(&structp.k) }},
	{"structp.x[0]", func() { println(structp.x[0]) }},
	{"&structp.x[0]", func() { println(&structp.x[0]) }},
	{"structp.x[i]", func() { println(structp.x[i]) }},
	{"&structp.x[i]", func() { println(&structp.x[i]) }},
	{"structp.x[9]", func() { println(structp.x[9]) }},
	{"&structp.x[9]", func() { println(&structp.x[9]) }},
	{"structp.l", func() { println(structp.l) }},
	{"&structp.l", func() { println(&structp.l) }},
	{"*bigstructp", func() { use(*bigstructp) }},
	{"&*bigstructp", func() { println(&*bigstructp) }},
	{"bigstructp.i", func() { println(bigstructp.i) }},
	{"&bigstructp.i", func() { println(&bigstructp.i) }},
	{"bigstructp.j", func() { println(bigstructp.j) }},
	{"&bigstructp.j", func() { println(&bigstructp.j) }},
	{"bigstructp.k", func() { println(bigstructp.k) }},
	{"&bigstructp.k", func() { println(&bigstructp.k) }},
	{"bigstructp.x[0]", func() { println(bigstructp.x[0]) }},
	{"&bigstructp.x[0]", func() { println(&bigstructp.x[0]) }},
	{"bigstructp.x[i]", func() { println(bigstructp.x[i]) }},
	{"&bigstructp.x[i]", func() { println(&bigstructp.x[i]) }},
	{"bigstructp.x[9]", func() { println(bigstructp.x[9]) }},
	{"&bigstructp.x[9]", func() { println(&bigstructp.x[9]) }},
	{"bigstructp.x[100<<20]", func() { println(bigstructp.x[100<<20]) }},
	{"&bigstructp.x[100<<20]", func() { println(&bigstructp.x[100<<20]) }},
	{"bigstructp.l", func() { println(bigstructp.l) }},
	{"&bigstructp.l", func() { println(&bigstructp.l) }},
	{"m1.F()", func() { println(m1.F()) }},
	{"m1.M.F()", func() { println(m1.M.F()) }},
	{"m2.F()", func() { println(m2.F()) }},
	{"m2.M.F()", func() { println(m2.M.F()) }},
}

type Struct struct {
	i int
	j float64
	k string
	x [10]int
	l []byte
}

type BigStruct struct {
	i int
	j float64
	k string
	x [128<<20]byte
	l []byte
}

type M struct {
}

func (m *M) F() int {return 0}

type M1 struct {
	M
}

type M2 struct {
	x int
	M
}
<|endoftext|>


<commit_before>package device

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/errors"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/subprocess"
	"github.com/lxc/lxd/shared/validate"
)

type tpm struct {
	deviceCommon
}

// CanMigrate returns whether the device can be migrated to any other cluster member.
func (d *tpm) CanMigrate() bool {
	return true
}

// validateConfig checks the supplied config for correctness.
func (d *tpm) validateConfig(instConf instance.ConfigReader) error {
	if !instanceSupported(instConf.Type(), instancetype.Container, instancetype.VM) {
		return ErrUnsupportedDevType
	}

	rules := map[string]func(string) error{
		"path": validate.IsNotEmpty,
	}

	err := d.config.Validate(rules)
	if err != nil {
		return errors.Wrapf(err, "Failed to validate config")
	}

	return nil
}

// validateEnvironment checks if the TPM emulator is available.
func (d *tpm) validateEnvironment() error {
	// Validate the required binary.
	_, err := exec.LookPath("swtpm")
	if err != nil {
		return fmt.Errorf("Required tool '%s' is missing", "swtpm")
	}

	if d.inst.Type() == instancetype.Container {
		// Load module tpm_vtpm_proxy which creates the /dev/vtpmx device, required
		// by the TPM emulator.
		module := "tpm_vtpm_proxy"

		err := util.LoadModule(module)
		if err != nil {
			return errors.Wrapf(err, "Failed to load kernel module %q", module)
		}
	}

	return nil
}

// Start is run when the device is added to the instance.
func (d *tpm) Start() (*deviceConfig.RunConfig, error) {
	err := d.validateEnvironment()
	if err != nil {
		return nil, errors.Wrap(err, "Failed to validate environment")
	}

	tpmDevPath := filepath.Join(d.inst.Path(), fmt.Sprintf("tpm.%s", d.name))

	if !shared.PathExists(tpmDevPath) {
		err := os.Mkdir(tpmDevPath, 0700)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to create device path %q", tpmDevPath)
		}
	}

	if d.inst.Type() == instancetype.VM {
		return d.startVM()
	}

	return d.startContainer()
}

func (d *tpm) startContainer() (*deviceConfig.RunConfig, error) {
	tpmDevPath := filepath.Join(d.inst.Path(), fmt.Sprintf("tpm.%s", d.name))
	logFileName := fmt.Sprintf("tpm.%s.log", d.name)
	logPath := filepath.Join(d.inst.LogPath(), logFileName)

	proc, err := subprocess.NewProcess("swtpm", []string{"chardev", "--tpm2", "--tpmstate", fmt.Sprintf("dir=%s", tpmDevPath), "--vtpm-proxy"}, logPath, "")
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to create new process")
	}

	err = proc.Start()
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to start process %q", "swtpm")
	}

	revert := revert.New()
	defer revert.Fail()

	// Stop the TPM emulator if anything goes wrong.
	revert.Add(func() { proc.Stop() })

	pidPath := filepath.Join(d.inst.DevicesPath(), fmt.Sprintf("%s.pid", d.name))

	err = proc.Save(pidPath)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to save swtpm state for device %q", d.name)
	}

	var major, minor int

	// We need to capture the output of the TPM emulator since it contains the device path. To do
	// that, we wait until something has been written to the log file (stdout redirect), and then
	// read it.
	for i := 0; i < 20; i++ {
		fi, err := os.Stat(logPath)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to stat %q", logPath)
		}

		if fi.Size() > 0 {
			break
		}

		time.Sleep(500 * time.Millisecond)
	}

	line, err := ioutil.ReadFile(logPath)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to read %q", logPath)
	}

	// The output will be something like:
	//   New TPM device: /dev/tpm1 (major/minor = 253/1)
	// We just need the major/minor numbers.
	fields := strings.Split(string(line), " ")

	if len(fields) < 7 {
		return nil, fmt.Errorf("Failed to get TPM device information")
	}

	_, err = fmt.Sscanf(fields[6], "%d/%d)", &major, &minor)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to retrieve major/minor number")
	}

	// Return error as we were unable to retrieve information regarding the TPM device.
	if major == 0 && minor == 0 {
		return nil, fmt.Errorf("Failed to get TPM device information")
	}

	runConf := deviceConfig.RunConfig{}

	err = unixDeviceSetupCharNum(d.state, d.inst.DevicesPath(), "unix", d.name, d.config, uint32(major), uint32(minor), d.config["path"], false, &runConf)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to setup unix device")
	}

	revert.Success()

	return &runConf, nil
}

func (d *tpm) startVM() (*deviceConfig.RunConfig, error) {
	tpmDevPath := filepath.Join(d.inst.Path(), fmt.Sprintf("tpm.%s", d.name))
	socketPath := filepath.Join(tpmDevPath, fmt.Sprintf("swtpm-%s.sock", d.name))
	runConf := deviceConfig.RunConfig{
		TPMDevice: []deviceConfig.RunConfigItem{
			{Key: "devName", Value: d.name},
			{Key: "path", Value: socketPath},
		},
	}

	proc, err := subprocess.NewProcess("swtpm", []string{"socket", "--tpm2", "--tpmstate", fmt.Sprintf("dir=%s", tpmDevPath), "--ctrl", fmt.Sprintf("type=unixio,path=%s", socketPath)}, "", "")
	if err != nil {
		return nil, err
	}

	// Start the TPM emulator.
	err = proc.Start()
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to start swtpm for device %q", d.name)
	}

	revert := revert.New()
	defer revert.Fail()

	revert.Add(func() { proc.Stop() })

	pidPath := filepath.Join(d.inst.DevicesPath(), fmt.Sprintf("%s.pid", d.name))

	err = proc.Save(pidPath)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to save swtpm state for device %q", d.name)
	}

	revert.Success()

	return &runConf, nil
}

// Stop terminates the TPM emulator.
func (d *tpm) Stop() (*deviceConfig.RunConfig, error) {
	pidPath := filepath.Join(d.inst.DevicesPath(), fmt.Sprintf("%s.pid", d.name))
	runConf := deviceConfig.RunConfig{}

	defer os.Remove(pidPath)

	if shared.PathExists(pidPath) {
		proc, err := subprocess.ImportProcess(pidPath)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to import process %q", pidPath)
		}

		// The TPM emulator will usually exit automatically when the tpm device is no longer in use,
		// i.e. the instance is stopped. Therefore, we only fail if the running process couldn't
		// be stopped.
		err = proc.Stop()
		if err != nil && err != subprocess.ErrNotRunning {
			return nil, errors.Wrapf(err, "Failed to stop imported process %q", pidPath)
		}
	}

	if d.inst.Type() == instancetype.Container {
		err := unixDeviceRemove(d.inst.DevicesPath(), "unix", d.name, "", &runConf)
		if err != nil {
			return nil, errors.Wrap(err, "Failed to remove unix device")
		}
	}

	return &runConf, nil
}

// Remove removes the TPM state file.
func (d *tpm) Remove() error {
	tpmDevPath := filepath.Join(d.inst.Path(), fmt.Sprintf("tpm.%s", d.name))

	return os.RemoveAll(tpmDevPath)
}
<commit_msg>lxd/device/tpm: Require `path` only for containers<commit_after>package device

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/errors"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/subprocess"
	"github.com/lxc/lxd/shared/validate"
)

type tpm struct {
	deviceCommon
}

// CanMigrate returns whether the device can be migrated to any other cluster member.
func (d *tpm) CanMigrate() bool {
	return true
}

// validateConfig checks the supplied config for correctness.
func (d *tpm) validateConfig(instConf instance.ConfigReader) error {
	if !instanceSupported(instConf.Type(), instancetype.Container, instancetype.VM) {
		return ErrUnsupportedDevType
	}

	rules := map[string]func(string) error{}

	if instConf.Type() == instancetype.Container {
		rules["path"] = validate.IsNotEmpty
	}

	err := d.config.Validate(rules)
	if err != nil {
		return errors.Wrapf(err, "Failed to validate config")
	}

	return nil
}

// validateEnvironment checks if the TPM emulator is available.
func (d *tpm) validateEnvironment() error {
	// Validate the required binary.
	_, err := exec.LookPath("swtpm")
	if err != nil {
		return fmt.Errorf("Required tool '%s' is missing", "swtpm")
	}

	if d.inst.Type() == instancetype.Container {
		// Load module tpm_vtpm_proxy which creates the /dev/vtpmx device, required
		// by the TPM emulator.
		module := "tpm_vtpm_proxy"

		err := util.LoadModule(module)
		if err != nil {
			return errors.Wrapf(err, "Failed to load kernel module %q", module)
		}
	}

	return nil
}

// Start is run when the device is added to the instance.
func (d *tpm) Start() (*deviceConfig.RunConfig, error) {
	err := d.validateEnvironment()
	if err != nil {
		return nil, errors.Wrap(err, "Failed to validate environment")
	}

	tpmDevPath := filepath.Join(d.inst.Path(), fmt.Sprintf("tpm.%s", d.name))

	if !shared.PathExists(tpmDevPath) {
		err := os.Mkdir(tpmDevPath, 0700)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to create device path %q", tpmDevPath)
		}
	}

	if d.inst.Type() == instancetype.VM {
		return d.startVM()
	}

	return d.startContainer()
}

func (d *tpm) startContainer() (*deviceConfig.RunConfig, error) {
	tpmDevPath := filepath.Join(d.inst.Path(), fmt.Sprintf("tpm.%s", d.name))
	logFileName := fmt.Sprintf("tpm.%s.log", d.name)
	logPath := filepath.Join(d.inst.LogPath(), logFileName)

	proc, err := subprocess.NewProcess("swtpm", []string{"chardev", "--tpm2", "--tpmstate", fmt.Sprintf("dir=%s", tpmDevPath), "--vtpm-proxy"}, logPath, "")
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to create new process")
	}

	err = proc.Start()
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to start process %q", "swtpm")
	}

	revert := revert.New()
	defer revert.Fail()

	// Stop the TPM emulator if anything goes wrong.
	revert.Add(func() { proc.Stop() })

	pidPath := filepath.Join(d.inst.DevicesPath(), fmt.Sprintf("%s.pid", d.name))

	err = proc.Save(pidPath)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to save swtpm state for device %q", d.name)
	}

	var major, minor int

	// We need to capture the output of the TPM emulator since it contains the device path. To do
	// that, we wait until something has been written to the log file (stdout redirect), and then
	// read it.
	for i := 0; i < 20; i++ {
		fi, err := os.Stat(logPath)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to stat %q", logPath)
		}

		if fi.Size() > 0 {
			break
		}

		time.Sleep(500 * time.Millisecond)
	}

	line, err := ioutil.ReadFile(logPath)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to read %q", logPath)
	}

	// The output will be something like:
	//   New TPM device: /dev/tpm1 (major/minor = 253/1)
	// We just need the major/minor numbers.
	fields := strings.Split(string(line), " ")

	if len(fields) < 7 {
		return nil, fmt.Errorf("Failed to get TPM device information")
	}

	_, err = fmt.Sscanf(fields[6], "%d/%d)", &major, &minor)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to retrieve major/minor number")
	}

	// Return error as we were unable to retrieve information regarding the TPM device.
	if major == 0 && minor == 0 {
		return nil, fmt.Errorf("Failed to get TPM device information")
	}

	runConf := deviceConfig.RunConfig{}

	err = unixDeviceSetupCharNum(d.state, d.inst.DevicesPath(), "unix", d.name, d.config, uint32(major), uint32(minor), d.config["path"], false, &runConf)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to setup unix device")
	}

	revert.Success()

	return &runConf, nil
}

func (d *tpm) startVM() (*deviceConfig.RunConfig, error) {
	tpmDevPath := filepath.Join(d.inst.Path(), fmt.Sprintf("tpm.%s", d.name))
	socketPath := filepath.Join(tpmDevPath, fmt.Sprintf("swtpm-%s.sock", d.name))
	runConf := deviceConfig.RunConfig{
		TPMDevice: []deviceConfig.RunConfigItem{
			{Key: "devName", Value: d.name},
			{Key: "path", Value: socketPath},
		},
	}

	proc, err := subprocess.NewProcess("swtpm", []string{"socket", "--tpm2", "--tpmstate", fmt.Sprintf("dir=%s", tpmDevPath), "--ctrl", fmt.Sprintf("type=unixio,path=%s", socketPath)}, "", "")
	if err != nil {
		return nil, err
	}

	// Start the TPM emulator.
	err = proc.Start()
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to start swtpm for device %q", d.name)
	}

	revert := revert.New()
	defer revert.Fail()

	revert.Add(func() { proc.Stop() })

	pidPath := filepath.Join(d.inst.DevicesPath(), fmt.Sprintf("%s.pid", d.name))

	err = proc.Save(pidPath)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to save swtpm state for device %q", d.name)
	}

	revert.Success()

	return &runConf, nil
}

// Stop terminates the TPM emulator.
func (d *tpm) Stop() (*deviceConfig.RunConfig, error) {
	pidPath := filepath.Join(d.inst.DevicesPath(), fmt.Sprintf("%s.pid", d.name))
	runConf := deviceConfig.RunConfig{}

	defer os.Remove(pidPath)

	if shared.PathExists(pidPath) {
		proc, err := subprocess.ImportProcess(pidPath)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to import process %q", pidPath)
		}

		// The TPM emulator will usually exit automatically when the tpm device is no longer in use,
		// i.e. the instance is stopped. Therefore, we only fail if the running process couldn't
		// be stopped.
		err = proc.Stop()
		if err != nil && err != subprocess.ErrNotRunning {
			return nil, errors.Wrapf(err, "Failed to stop imported process %q", pidPath)
		}
	}

	if d.inst.Type() == instancetype.Container {
		err := unixDeviceRemove(d.inst.DevicesPath(), "unix", d.name, "", &runConf)
		if err != nil {
			return nil, errors.Wrap(err, "Failed to remove unix device")
		}
	}

	return &runConf, nil
}

// Remove removes the TPM state file.
func (d *tpm) Remove() error {
	tpmDevPath := filepath.Join(d.inst.Path(), fmt.Sprintf("tpm.%s", d.name))

	return os.RemoveAll(tpmDevPath)
}
<|endoftext|>


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

// +build dragonfly freebsd js,wasm linux nacl netbsd openbsd solaris windows

package poll

import "syscall"

// Fsync wraps syscall.Fsync.
func (fd *FD) Fsync() error {
	if err := fd.incref(); err != nil {
		return err
	}
	defer fd.decref()
	return syscall.Fsync(fd.Sysfd)
}
<commit_msg>internal/poll: add FD.Fsync on aix<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build aix dragonfly freebsd js,wasm linux nacl netbsd openbsd solaris windows

package poll

import "syscall"

// Fsync wraps syscall.Fsync.
func (fd *FD) Fsync() error {
	if err := fd.incref(); err != nil {
		return err
	}
	defer fd.decref()
	return syscall.Fsync(fd.Sysfd)
}
<|endoftext|>


<commit_before><commit_msg>Key function description map by SSA fn name, not type object<commit_after><|endoftext|>


<commit_before>package middleware

import (
	"os"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/gin-gonic/gin"
)

func Logging() gin.HandlerFunc {
	return func(c *gin.Context) {
		logger := logrus.StandardLogger()
		start := time.Now().UTC()
		path := c.Request.URL.Path
		c.Next()
		end := time.Now().UTC()
		latency := end.Sub(start)

		formatter := &logrus.TextFormatter{
			FullTimestamp:   true,
			TimestampFormat: "2006-01-02-15:04:05",
		}
		logrus.SetFormatter(formatter)
		//logrus.SetFormatter(&logrus.JSONFormatter{})
		file, err := os.OpenFile("./logrus.log", os.O_APPEND|os.O_WRONLY, 0666)
		if err == nil {
			logger.Out = file
		} else {
			logger.Info("Failed to log to file, using default stderr")
		}

		logger.WithFields(logrus.Fields{
			"status":     c.Writer.Status(),
			"method":     c.Request.Method,
			"path":       path,
			"ip":         c.ClientIP(),
			"duration":   latency,
			"user_agent": c.Request.UserAgent(),
		}).Info()
	}
}
<commit_msg>New format for route logging<commit_after>package middleware

import (
	"fmt"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/TeaMeow/KitSvc/module/logger"
	"github.com/gin-gonic/gin"
	"github.com/willf/pad"
)

func Logging() gin.HandlerFunc {
	return func(c *gin.Context) {

		start := time.Now().UTC()
		path := c.Request.URL.Path
		c.Next()

		if path == "/metrics" || path == "/sd/health" || path == "/sd/ram" || path == "/sd/cpu" || path == "/sd/disk" {
			return
		}

		end := time.Now().UTC()
		latency := end.Sub(start)

		status := c.Writer.Status()
		method := c.Request.Method
		ip := c.ClientIP()
		userAgent := c.Request.UserAgent()

		logger.InfoFields(fmt.Sprintf("%d | %s | %s | %s %s",
			status,
			pad.Right(latency.String(), 13, " "),
			pad.Right(ip, 12, " "),
			pad.Right(method, 5, " "),
			pad.Right(path, 15, " "),
		), logrus.Fields{
			//"status":     status,
			//"method":     method,
			//"path":       path,
			//"ip":         ip,
			//"duration":   latency,
			"user_agent": userAgent,
		})
	}
}
<|endoftext|>


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

// HTTP reverse proxy handler

package proxy

import (
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration

	// ErrorLog specifies an optional logger for errors
	// that occur when attempting to proxy the request.
	// If nil, logging goes to os.Stderr via the log package's
	// standard logger.
	ErrorLog *log.Logger
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	outreq.URL.Scheme = "http"
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		p.logf("http: proxy error: %v", err)
		rw.WriteHeader(http.StatusInternalServerError)
		return
	}
	defer res.Body.Close()

	for _, h := range hopHeaders {
		res.Header.Del(h)
	}

	copyHeader(rw.Header(), res.Header)

	rw.WriteHeader(res.StatusCode)
	p.copyResponse(rw, res.Body)
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	io.Copy(dst, src)
}

func (p *ReverseProxy) logf(format string, args ...interface{}) {
	if p.ErrorLog != nil {
		p.ErrorLog.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<commit_msg>router: refactor proxy request preparation & response writing<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// HTTP reverse proxy handler

package proxy

import (
	"io"
	"log"
	"net/http"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration

	// ErrorLog specifies an optional logger for errors
	// that occur when attempting to proxy the request.
	// If nil, logging goes to os.Stderr via the log package's
	// standard logger.
	ErrorLog *log.Logger
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := prepareRequest(req)

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		p.logf("router: proxy error: %v", err)
		rw.WriteHeader(http.StatusServiceUnavailable)
		rw.Write([]byte(http.StatusText(http.StatusServiceUnavailable)))
		return
	}
	defer res.Body.Close()

	p.writeResponse(rw, res)
}

func (p *ReverseProxy) writeResponse(rw http.ResponseWriter, res *http.Response) {
	for _, h := range hopHeaders {
		res.Header.Del(h)
	}

	copyHeader(rw.Header(), res.Header)

	rw.WriteHeader(res.StatusCode)
	p.copyResponse(rw, res.Body)
}

func prepareRequest(req *http.Request) *http.Request {
	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	outreq.URL.Scheme = "http"
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	return outreq
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	io.Copy(dst, src)
}

func (p *ReverseProxy) logf(format string, args ...interface{}) {
	if p.ErrorLog != nil {
		p.ErrorLog.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<|endoftext|>