<commit_before>/*
Copyright 2019 The Skaffold 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 util

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strings"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

const (
	hiddenPrefix string = "."
)

func RandomID() string {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", b)
}

func RandomFourCharacterID() string {
	b := make([]byte, 2)
	_, err := rand.Read(b)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", b)
}

// These are the supported file formats for kubernetes manifests
var validSuffixes = []string{".yml", ".yaml", ".json"}

// IsSupportedKubernetesFormat is for determining if a file under a glob pattern
// is deployable file format. It makes no attempt to check whether or not the file
// is actually deployable or has the correct contents.
func IsSupportedKubernetesFormat(n string) bool {
	for _, s := range validSuffixes {
		if strings.HasSuffix(n, s) {
			return true
		}
	}
	return false
}

func StrSliceContains(sl []string, s string) bool {
	for _, a := range sl {
		if a == s {
			return true
		}
	}
	return false
}

// ExpandPathsGlob expands paths according to filepath.Glob patterns
// Returns a list of unique files that match the glob patterns passed in.
func ExpandPathsGlob(workingDir string, paths []string) ([]string, error) {
	expandedPaths := make(map[string]bool)
	for _, p := range paths {
		if filepath.IsAbs(p) {
			// This is a absolute file reference
			expandedPaths[p] = true
			continue
		}

		path := filepath.Join(workingDir, p)

		if _, err := os.Stat(path); err == nil {
			// This is a file reference, so just add it
			expandedPaths[path] = true
			continue
		}

		files, err := filepath.Glob(path)
		if err != nil {
			return nil, errors.Wrap(err, "glob")
		}
		if len(files) == 0 {
			logrus.Warnf("%s did not match any file", p)
		}

		for _, f := range files {
			err := filepath.Walk(f, func(path string, info os.FileInfo, err error) error {
				if !info.IsDir() {
					expandedPaths[path] = true
				}

				return nil
			})
			if err != nil {
				return nil, errors.Wrap(err, "filepath walk")
			}
		}
	}

	var ret []string
	for k := range expandedPaths {
		ret = append(ret, k)
	}
	sort.Strings(ret)
	return ret, nil
}

// HasMeta reports whether path contains any of the magic characters
// recognized by filepath.Match.
// This is a copy of filepath/match.go's hasMeta
func HasMeta(path string) bool {
	magicChars := `*?[`
	if runtime.GOOS != "windows" {
		magicChars = `*?[\`
	}
	return strings.ContainsAny(path, magicChars)
}

// BoolPtr returns a pointer to a bool
func BoolPtr(b bool) *bool {
	o := b
	return &o
}

// StringPtr returns a pointer to a string
func StringPtr(s string) *string {
	o := s
	return &o
}

func IsURL(s string) bool {
	return strings.HasPrefix(s, "http://") || strings.HasPrefix(s, "https://")
}

func Download(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}

// VerifyOrCreateFile checks if a file exists at the given path,
// and if not, creates all parent directories and creates the file.
func VerifyOrCreateFile(path string) error {
	_, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		dir := filepath.Dir(path)
		if err = os.MkdirAll(dir, 0744); err != nil {
			return errors.Wrap(err, "creating parent directory")
		}
		if _, err = os.Create(path); err != nil {
			return errors.Wrap(err, "creating file")
		}
		return nil
	}
	return err
}

// RemoveFromSlice removes a string from a slice of strings
func RemoveFromSlice(s []string, target string) []string {
	for i, val := range s {
		if val == target {
			return append(s[:i], s[i+1:]...)
		}
	}
	return s
}

// Expand replaces placeholders for a given key with a given value.
// It supports the ${key} and the $key syntax.
func Expand(text, key, value string) string {
	text = strings.Replace(text, "${"+key+"}", value, -1)

	indices := regexp.MustCompile(`\$`+key).FindAllStringIndex(text, -1)

	for i := len(indices) - 1; i >= 0; i-- {
		from := indices[i][0]
		to := indices[i][1]

		if to >= len(text) || !isAlphaNum(text[to]) {
			text = text[0:from] + value + text[to:]
		}
	}

	return text
}

func isAlphaNum(c uint8) bool {
	return c == '_' || '0' <= c && c <= '9' || 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z'
}

// AbsFile resolves the absolute path of the file named filename in directory workspace, erroring if it is not a file
func AbsFile(workspace string, filename string) (string, error) {
	file := filepath.Join(workspace, filename)
	info, err := os.Stat(file)
	if err != nil {
		return "", err
	}
	if info.IsDir() {
		return "", errors.Errorf("%s is a directory", file)
	}
	return filepath.Abs(file)
}

// NonEmptyLines scans the provided input and returns the non-empty strings found as an array
func NonEmptyLines(input []byte) []string {
	var result []string
	scanner := bufio.NewScanner(bytes.NewReader(input))
	for scanner.Scan() {
		if line := scanner.Text(); len(line) > 0 {
			result = append(result, line)
		}
	}
	return result
}

// SHA256 returns the shasum of the contents of r
func SHA256(r io.Reader) (string, error) {
	hasher := sha256.New()
	_, err := io.Copy(hasher, r)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(hasher.Sum(make([]byte, 0, hasher.Size()))), nil
}

// CloneThroughJSON marshals the old interface into the new one
func CloneThroughJSON(old interface{}, new interface{}) error {
	o, err := json.Marshal(old)
	if err != nil {
		return errors.Wrap(err, "marshalling old")
	}
	if err := json.Unmarshal(o, &new); err != nil {
		return errors.Wrap(err, "unmarshalling new")
	}
	return nil
}

// AbsolutePaths prepends each path in paths with workspace if the path isn't absolute
func AbsolutePaths(workspace string, paths []string) []string {
	var p []string
	for _, path := range paths {
		// TODO(dgageot): this is only done for jib builder.
		if !filepath.IsAbs(path) {
			path = filepath.Join(workspace, path)
		}
		p = append(p, path)
	}
	return p
}

// IsHiddenDir returns if a directory is hidden.
func IsHiddenDir(filename string) bool {
	// Return false for current dir
	if filename == hiddenPrefix {
		return false
	}
	return hasHiddenPrefix(filename)
}

// IsHiddenFile returns if a file is hidden.
// File is hidden if it starts with prefix "."
func IsHiddenFile(filename string) bool {
	return hasHiddenPrefix(filename)
}

func hasHiddenPrefix(s string) bool {
	return strings.HasPrefix(s, hiddenPrefix)
}
<commit_msg>Remove dead code<commit_after>/*
Copyright 2019 The Skaffold 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 util

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

const (
	hiddenPrefix string = "."
)

func RandomID() string {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", b)
}

// These are the supported file formats for kubernetes manifests
var validSuffixes = []string{".yml", ".yaml", ".json"}

// IsSupportedKubernetesFormat is for determining if a file under a glob pattern
// is deployable file format. It makes no attempt to check whether or not the file
// is actually deployable or has the correct contents.
func IsSupportedKubernetesFormat(n string) bool {
	for _, s := range validSuffixes {
		if strings.HasSuffix(n, s) {
			return true
		}
	}
	return false
}

func StrSliceContains(sl []string, s string) bool {
	for _, a := range sl {
		if a == s {
			return true
		}
	}
	return false
}

// ExpandPathsGlob expands paths according to filepath.Glob patterns
// Returns a list of unique files that match the glob patterns passed in.
func ExpandPathsGlob(workingDir string, paths []string) ([]string, error) {
	expandedPaths := make(map[string]bool)
	for _, p := range paths {
		if filepath.IsAbs(p) {
			// This is a absolute file reference
			expandedPaths[p] = true
			continue
		}

		path := filepath.Join(workingDir, p)

		if _, err := os.Stat(path); err == nil {
			// This is a file reference, so just add it
			expandedPaths[path] = true
			continue
		}

		files, err := filepath.Glob(path)
		if err != nil {
			return nil, errors.Wrap(err, "glob")
		}
		if len(files) == 0 {
			logrus.Warnf("%s did not match any file", p)
		}

		for _, f := range files {
			err := filepath.Walk(f, func(path string, info os.FileInfo, err error) error {
				if !info.IsDir() {
					expandedPaths[path] = true
				}

				return nil
			})
			if err != nil {
				return nil, errors.Wrap(err, "filepath walk")
			}
		}
	}

	var ret []string
	for k := range expandedPaths {
		ret = append(ret, k)
	}
	sort.Strings(ret)
	return ret, nil
}

// BoolPtr returns a pointer to a bool
func BoolPtr(b bool) *bool {
	o := b
	return &o
}

// StringPtr returns a pointer to a string
func StringPtr(s string) *string {
	o := s
	return &o
}

func IsURL(s string) bool {
	return strings.HasPrefix(s, "http://") || strings.HasPrefix(s, "https://")
}

func Download(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}

// VerifyOrCreateFile checks if a file exists at the given path,
// and if not, creates all parent directories and creates the file.
func VerifyOrCreateFile(path string) error {
	_, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		dir := filepath.Dir(path)
		if err = os.MkdirAll(dir, 0744); err != nil {
			return errors.Wrap(err, "creating parent directory")
		}
		if _, err = os.Create(path); err != nil {
			return errors.Wrap(err, "creating file")
		}
		return nil
	}
	return err
}

// RemoveFromSlice removes a string from a slice of strings
func RemoveFromSlice(s []string, target string) []string {
	for i, val := range s {
		if val == target {
			return append(s[:i], s[i+1:]...)
		}
	}
	return s
}

// Expand replaces placeholders for a given key with a given value.
// It supports the ${key} and the $key syntax.
func Expand(text, key, value string) string {
	text = strings.Replace(text, "${"+key+"}", value, -1)

	indices := regexp.MustCompile(`\$`+key).FindAllStringIndex(text, -1)

	for i := len(indices) - 1; i >= 0; i-- {
		from := indices[i][0]
		to := indices[i][1]

		if to >= len(text) || !isAlphaNum(text[to]) {
			text = text[0:from] + value + text[to:]
		}
	}

	return text
}

func isAlphaNum(c uint8) bool {
	return c == '_' || '0' <= c && c <= '9' || 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z'
}

// AbsFile resolves the absolute path of the file named filename in directory workspace, erroring if it is not a file
func AbsFile(workspace string, filename string) (string, error) {
	file := filepath.Join(workspace, filename)
	info, err := os.Stat(file)
	if err != nil {
		return "", err
	}
	if info.IsDir() {
		return "", errors.Errorf("%s is a directory", file)
	}
	return filepath.Abs(file)
}

// NonEmptyLines scans the provided input and returns the non-empty strings found as an array
func NonEmptyLines(input []byte) []string {
	var result []string
	scanner := bufio.NewScanner(bytes.NewReader(input))
	for scanner.Scan() {
		if line := scanner.Text(); len(line) > 0 {
			result = append(result, line)
		}
	}
	return result
}

// SHA256 returns the shasum of the contents of r
func SHA256(r io.Reader) (string, error) {
	hasher := sha256.New()
	_, err := io.Copy(hasher, r)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(hasher.Sum(make([]byte, 0, hasher.Size()))), nil
}

// CloneThroughJSON marshals the old interface into the new one
func CloneThroughJSON(old interface{}, new interface{}) error {
	o, err := json.Marshal(old)
	if err != nil {
		return errors.Wrap(err, "marshalling old")
	}
	if err := json.Unmarshal(o, &new); err != nil {
		return errors.Wrap(err, "unmarshalling new")
	}
	return nil
}

// AbsolutePaths prepends each path in paths with workspace if the path isn't absolute
func AbsolutePaths(workspace string, paths []string) []string {
	var p []string
	for _, path := range paths {
		// TODO(dgageot): this is only done for jib builder.
		if !filepath.IsAbs(path) {
			path = filepath.Join(workspace, path)
		}
		p = append(p, path)
	}
	return p
}

// IsHiddenDir returns if a directory is hidden.
func IsHiddenDir(filename string) bool {
	// Return false for current dir
	if filename == hiddenPrefix {
		return false
	}
	return hasHiddenPrefix(filename)
}

// IsHiddenFile returns if a file is hidden.
// File is hidden if it starts with prefix "."
func IsHiddenFile(filename string) bool {
	return hasHiddenPrefix(filename)
}

func hasHiddenPrefix(s string) bool {
	return strings.HasPrefix(s, hiddenPrefix)
}
<|endoftext|>


<commit_before>package runner

import (
	"context"
	"errors"
	"fmt"
	"sync"

	"io"
	"time"

	core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
	cds "github.com/envoyproxy/go-control-plane/envoy/service/cluster/v3"
	discovery "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"

	"github.com/ii/xds-test-harness/internal/parser"
	pb "github.com/ii/xds-test-harness/api/adapter"
	"google.golang.org/grpc"
	"github.com/rs/zerolog/log"
)

var (
	opts []grpc.DialOption = []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithBlock(),
		grpc.WithTimeout(time.Second * 10),
	}
)

type ClientConfig struct {
	Port string
	Conn *grpc.ClientConn
}

type Cache struct {
	StartState *pb.Snapshot
	FinalResponse *discovery.DiscoveryResponse
}

type CDS struct {
	Req   chan *discovery.DiscoveryRequest
	Res   chan *discovery.DiscoveryResponse
	Err   chan error
	Done  chan bool
	Cache struct {
		InitResource []string
		Requests  []*discovery.DiscoveryRequest
		Responses []*discovery.DiscoveryResponse
	}
}

type Runner struct {
	Adapter *ClientConfig
	Target  *ClientConfig
	NodeID  string
	Cache   *Cache
	CDS     *CDS
}

func NewRunner() *Runner {
	return &Runner{
		Adapter: &ClientConfig{},
		Target:  &ClientConfig{},
		NodeID:  "",
		Cache:   &Cache{},
		CDS:     &CDS{},
	}
}

func FreshRunner (current *Runner) *Runner {
	return &Runner{
		Adapter: current.Adapter,
		Target: current.Target,
		NodeID: current.NodeID,
		Cache: &Cache{},
		CDS: &CDS{},
	}
}

func connectViaGRPC(client *ClientConfig, server string) (conn *grpc.ClientConn, err error) {
	conn, err = grpc.Dial(client.Port, opts...)
	if err != nil {
		err = fmt.Errorf("Cannot connect at %v: %v\n", client.Port, err)
		return nil, err
	}
	log.Debug().
		Msgf("Runner connected to %v", server)
	return conn, nil
}

func (r *Runner) ConnectToTarget(address string) error {
	r.Target.Port = address
	conn, err := connectViaGRPC(r.Target, "target")
	if err != nil {
		return err
	}
	r.Target.Conn = conn
	return nil
}

func (r *Runner) ConnectToAdapter(address string) error {
	r.Adapter.Port = address
	conn, err := connectViaGRPC(r.Adapter, "adapter")
	if err != nil {
		return err
	}
	r.Adapter.Conn = conn
	return nil
}

func (r *Runner) NewCDSRequest(resourceList []string) *discovery.DiscoveryRequest {
	clusters := []string{}
	for _, cluster := range resourceList {
		clusters = append(clusters, cluster)
	}
	return &discovery.DiscoveryRequest{
		VersionInfo: "",
		Node: &core.Node{
			Id: r.NodeID,
		},
		ResourceNames: clusters,
		TypeUrl:       "type.googleapis.com/envoy.config.cluster.v3.Cluster",
	}
}

func (r *Runner) NewCDSAckFromResponse(res *discovery.DiscoveryResponse) (*discovery.DiscoveryRequest, error) {
	response, err := parser.ParseDiscoveryResponse(res)
	if err != nil {
		err := fmt.Errorf("error parsing dres for acking: %v", err)
		return nil, err
	}

	request := &discovery.DiscoveryRequest{
		VersionInfo:   response.VersionInfo,
		ResourceNames: r.CDS.Cache.InitResource,
		TypeUrl:       "type.googleapis.com/envoy.config.cluster.v3.Cluster",
		ResponseNonce: response.Nonce,
	}
	return request, nil
}

func (r *Runner) AckCDS(initReq *discovery.DiscoveryRequest) {

	log.Debug().Msgf("Sending First Discovery Request", initReq)
	r.CDS.Req <- initReq
	r.CDS.Cache.Requests = append(r.CDS.Cache.Requests, initReq)

	for {
		select {
		case res := <-r.CDS.Res:
			r.CDS.Cache.Responses = append(r.CDS.Cache.Responses, res)
			ack, err := r.NewCDSAckFromResponse(res)
			if err != nil {
				log.Err(err).Msg("Error creating Ack Request")
			}
			log.Debug().
				Msgf("Sending Ack: %v", ack)
			r.CDS.Req <- ack
	        r.CDS.Cache.Requests = append(r.CDS.Cache.Requests, ack)
		case <-r.CDS.Done:
			log.Debug().Msg("Received Done signal, shutting down request channel")
			close(r.CDS.Req)
			return
		}
	}
}

func (r *Runner) CDSStream() error {

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	defer close(r.CDS.Err)

	client := cds.NewClusterDiscoveryServiceClient(r.Target.Conn)
	stream, err := client.StreamClusters(ctx)
	if err != nil {
		err = errors.New(fmt.Sprintf("Cannot start CDS stream %v. error: %v", stream, err))
		log.Debug().
			Err(err).
			Msg("")
		r.CDS.Err <- err
	}
	var wg sync.WaitGroup
	go func() {
		for {
			wg.Add(1)
			in, err := stream.Recv()
			if err == io.EOF {
				log.Debug().
					Msg("No more Discovery Responses from CDS stream")
				close(r.CDS.Res)
				return
			}
			if err != nil {
				log.Err(err).Msg("error receiving responses on CDS stream")
				r.CDS.Err <- err
			}
			log.Debug().
				Msgf("Received discovery response: %v", in)
			r.CDS.Res <- in
		}
	}()

	for req := range r.CDS.Req {
		if err := stream.Send(req); err != nil {
			log.Err(err).
				Msg("Error sending discovery request")
			r.CDS.Err <- err
		}
	}
	stream.CloseSend()
	wg.Wait()
	return err
}
<commit_msg>log error as error not debug<commit_after>package runner

import (
	"context"
	"errors"
	"fmt"
	"sync"

	"io"
	"time"

	core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
	cds "github.com/envoyproxy/go-control-plane/envoy/service/cluster/v3"
	discovery "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"

	"github.com/ii/xds-test-harness/internal/parser"
	pb "github.com/ii/xds-test-harness/api/adapter"
	"google.golang.org/grpc"
	"github.com/rs/zerolog/log"
)

var (
	opts []grpc.DialOption = []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithBlock(),
		grpc.WithTimeout(time.Second * 10),
	}
)

type ClientConfig struct {
	Port string
	Conn *grpc.ClientConn
}

type Cache struct {
	StartState *pb.Snapshot
	FinalResponse *discovery.DiscoveryResponse
}

type CDS struct {
	Req   chan *discovery.DiscoveryRequest
	Res   chan *discovery.DiscoveryResponse
	Err   chan error
	Done  chan bool
	Cache struct {
		InitResource []string
		Requests  []*discovery.DiscoveryRequest
		Responses []*discovery.DiscoveryResponse
	}
}

type Runner struct {
	Adapter *ClientConfig
	Target  *ClientConfig
	NodeID  string
	Cache   *Cache
	CDS     *CDS
}

func NewRunner() *Runner {
	return &Runner{
		Adapter: &ClientConfig{},
		Target:  &ClientConfig{},
		NodeID:  "",
		Cache:   &Cache{},
		CDS:     &CDS{},
	}
}

func FreshRunner (current *Runner) *Runner {
	return &Runner{
		Adapter: current.Adapter,
		Target: current.Target,
		NodeID: current.NodeID,
		Cache: &Cache{},
		CDS: &CDS{},
	}
}

func connectViaGRPC(client *ClientConfig, server string) (conn *grpc.ClientConn, err error) {
	conn, err = grpc.Dial(client.Port, opts...)
	if err != nil {
		err = fmt.Errorf("Cannot connect at %v: %v\n", client.Port, err)
		return nil, err
	}
	log.Debug().
		Msgf("Runner connected to %v", server)
	return conn, nil
}

func (r *Runner) ConnectToTarget(address string) error {
	r.Target.Port = address
	conn, err := connectViaGRPC(r.Target, "target")
	if err != nil {
		return err
	}
	r.Target.Conn = conn
	return nil
}

func (r *Runner) ConnectToAdapter(address string) error {
	r.Adapter.Port = address
	conn, err := connectViaGRPC(r.Adapter, "adapter")
	if err != nil {
		return err
	}
	r.Adapter.Conn = conn
	return nil
}

func (r *Runner) NewCDSRequest(resourceList []string) *discovery.DiscoveryRequest {
	clusters := []string{}
	for _, cluster := range resourceList {
		clusters = append(clusters, cluster)
	}
	return &discovery.DiscoveryRequest{
		VersionInfo: "",
		Node: &core.Node{
			Id: r.NodeID,
		},
		ResourceNames: clusters,
		TypeUrl:       "type.googleapis.com/envoy.config.cluster.v3.Cluster",
	}
}

func (r *Runner) NewCDSAckFromResponse(res *discovery.DiscoveryResponse) (*discovery.DiscoveryRequest, error) {
	response, err := parser.ParseDiscoveryResponse(res)
	if err != nil {
		err := fmt.Errorf("error parsing dres for acking: %v", err)
		return nil, err
	}

	request := &discovery.DiscoveryRequest{
		VersionInfo:   response.VersionInfo,
		ResourceNames: r.CDS.Cache.InitResource,
		TypeUrl:       "type.googleapis.com/envoy.config.cluster.v3.Cluster",
		ResponseNonce: response.Nonce,
	}
	return request, nil
}

func (r *Runner) AckCDS(initReq *discovery.DiscoveryRequest) {

	log.Debug().Msgf("Sending First Discovery Request", initReq)
	r.CDS.Req <- initReq
	r.CDS.Cache.Requests = append(r.CDS.Cache.Requests, initReq)

	for {
		select {
		case res := <-r.CDS.Res:
			r.CDS.Cache.Responses = append(r.CDS.Cache.Responses, res)
			ack, err := r.NewCDSAckFromResponse(res)
			if err != nil {
				log.Err(err).Msg("Error creating Ack Request")
			}
			log.Debug().
				Msgf("Sending Ack: %v", ack)
			r.CDS.Req <- ack
	        r.CDS.Cache.Requests = append(r.CDS.Cache.Requests, ack)
		case <-r.CDS.Done:
			log.Debug().Msg("Received Done signal, shutting down request channel")
			close(r.CDS.Req)
			return
		}
	}
}

func (r *Runner) CDSStream() error {

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	defer close(r.CDS.Err)

	client := cds.NewClusterDiscoveryServiceClient(r.Target.Conn)
	stream, err := client.StreamClusters(ctx)
	if err != nil {
		err = errors.New(fmt.Sprintf("Cannot start CDS stream %v. error: %v", stream, err))
		log.Debug().
			Err(err).
			Msg("")
		r.CDS.Err <- err
	}
	var wg sync.WaitGroup
	go func() {
		for {
			wg.Add(1)
			in, err := stream.Recv()
			if err == io.EOF {
				log.Debug().
					Msg("No more Discovery Responses from CDS stream")
				close(r.CDS.Res)
				return
			}
			if err != nil {
				log.Err(err).Msg("error receiving responses on CDS stream")
				r.CDS.Err <- err
				return
			}
			log.Debug().
				Msgf("Received discovery response: %v", in)
			r.CDS.Res <- in
		}
	}()

	for req := range r.CDS.Req {
		if err := stream.Send(req); err != nil {
			log.Err(err).
				Msg("Error sending discovery request")
			r.CDS.Err <- err
		}
	}
	stream.CloseSend()
	wg.Wait()
	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Gravitational, Inc.

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 service

import (
	"context"
	"sync"

	"github.com/gravitational/teleport/lib/utils"

	"github.com/gravitational/trace"
	"github.com/sirupsen/logrus"
)

// Supervisor implements the simple service logic - registering
// service functions and de-registering the service goroutines
type Supervisor interface {
	// Register adds the service to the pool, if supervisor is in
	// the started state, the service will be started immediately
	// otherwise, it will be started after Start() has been called
	Register(srv Service)

	// RegisterFunc creates a service from function spec and registers
	// it within the system
	RegisterFunc(name string, fn ServiceFunc)

	// ServiceCount returns the number of registered and actively running
	// services
	ServiceCount() int

	// Start starts all unstarted services
	Start() error

	// Wait waits until all services exit
	Wait() error

	// Run starts and waits for the service to complete
	// it's a combinatioin Start() and Wait()
	Run() error

	// Services returns list of running services
	Services() []string

	// BroadcastEvent generates event and broadcasts it to all
	// interested parties
	BroadcastEvent(Event)

	// WaitForEvent waits for event to be broadcasted, if the event
	// was already broadcasted, payloadC will receive current event immediately
	// CLose 'cancelC' channel to force WaitForEvent to return prematurely
	WaitForEvent(name string, eventC chan Event, cancelC chan struct{})
}

type LocalSupervisor struct {
	state int
	sync.Mutex
	wg           *sync.WaitGroup
	services     []Service
	errors       []error
	events       map[string]Event
	eventsC      chan Event
	eventWaiters map[string][]*waiter
	closeContext context.Context
	signalClose  context.CancelFunc
}

// NewSupervisor returns new instance of initialized supervisor
func NewSupervisor() Supervisor {
	closeContext, cancel := context.WithCancel(context.TODO())
	srv := &LocalSupervisor{
		services:     []Service{},
		wg:           &sync.WaitGroup{},
		events:       map[string]Event{},
		eventsC:      make(chan Event, 1024),
		eventWaiters: make(map[string][]*waiter),
		closeContext: closeContext,
		signalClose:  cancel,
	}
	go srv.fanOut()
	return srv
}

// Event is a special service event that can be generated
// by various goroutines in the supervisor
type Event struct {
	Name    string
	Payload interface{}
}

func (e *Event) String() string {
	return e.Name
}

func (s *LocalSupervisor) Register(srv Service) {
	log.WithFields(logrus.Fields{"service": srv.Name()}).Debugf("Adding service to supervisor.")
	s.Lock()
	defer s.Unlock()
	s.services = append(s.services, srv)

	if s.state == stateStarted {
		s.serve(srv)
	}
}

// ServiceCount returns the number of registered and actively running services
func (s *LocalSupervisor) ServiceCount() int {
	s.Lock()
	defer s.Unlock()
	return len(s.services)
}

// RegisterFunc creates a service from function spec and registers
// it within the system
func (s *LocalSupervisor) RegisterFunc(name string, fn ServiceFunc) {
	s.Register(&LocalService{Function: fn, ServiceName: name})
}

// RemoveService removes service from supervisor tracking list
func (s *LocalSupervisor) RemoveService(srv Service) error {
	log := log.WithFields(logrus.Fields{"service": srv.Name()})
	s.Lock()
	defer s.Unlock()
	for i, el := range s.services {
		if el == srv {
			s.services = append(s.services[:i], s.services[i+1:]...)
			log.Debugf("Service is completed and removed.")
			return nil
		}
	}
	log.Warningf("Service is completed but not found.")
	return trace.NotFound("service %v is not found", srv)
}

func (s *LocalSupervisor) serve(srv Service) {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		defer s.RemoveService(srv)
		log.WithFields(logrus.Fields{"service": srv.Name()}).Debugf("Service has started.")
		err := srv.Serve()
		if err != nil {
			utils.FatalError(err)
		}
	}()
}

func (s *LocalSupervisor) Start() error {
	s.Lock()
	defer s.Unlock()
	s.state = stateStarted

	if len(s.services) == 0 {
		log.Warning("Supervisor has no services to run. Exiting.")
		return nil
	}

	for _, srv := range s.services {
		s.serve(srv)
	}

	return nil
}

func (s *LocalSupervisor) Services() []string {
	s.Lock()
	defer s.Unlock()

	out := make([]string, len(s.services))

	for i, srv := range s.services {
		out[i] = srv.Name()
	}
	return out
}

func (s *LocalSupervisor) Wait() error {
	defer s.signalClose()
	s.wg.Wait()
	return nil
}

func (s *LocalSupervisor) Run() error {
	if err := s.Start(); err != nil {
		return trace.Wrap(err)
	}
	return s.Wait()
}

func (s *LocalSupervisor) BroadcastEvent(event Event) {
	s.Lock()
	defer s.Unlock()
	s.events[event.Name] = event
	log.WithFields(logrus.Fields{"event": event.String()}).Debugf("Broadcasting event.")

	go func() {
		s.eventsC <- event
	}()
}

func (s *LocalSupervisor) WaitForEvent(name string, eventC chan Event, cancelC chan struct{}) {
	s.Lock()
	defer s.Unlock()

	waiter := &waiter{eventC: eventC, cancelC: cancelC}
	event, ok := s.events[name]
	if ok {
		go s.notifyWaiter(waiter, event)
		return
	}
	s.eventWaiters[name] = append(s.eventWaiters[name], waiter)
}

func (s *LocalSupervisor) getWaiters(name string) []*waiter {
	s.Lock()
	defer s.Unlock()

	waiters := s.eventWaiters[name]
	out := make([]*waiter, len(waiters))
	for i := range waiters {
		out[i] = waiters[i]
	}
	return out
}

func (s *LocalSupervisor) notifyWaiter(w *waiter, event Event) {
	select {
	case w.eventC <- event:
	case <-w.cancelC:
	}
}

func (s *LocalSupervisor) fanOut() {
	for {
		select {
		case event := <-s.eventsC:
			waiters := s.getWaiters(event.Name)
			for _, waiter := range waiters {
				go s.notifyWaiter(waiter, event)
			}
		case <-s.closeContext.Done():
			return
		}
	}
}

type waiter struct {
	eventC  chan Event
	cancelC chan struct{}
}

// Service is a running teleport service function
type Service interface {
	// Serve starts the function
	Serve() error
	// String returns user-friendly description of service
	String() string
	// Name returns service name
	Name() string
}

// LocalService is a locally defined service
type LocalService struct {
	// Function is a function to call
	Function ServiceFunc
	// ServiceName is a service name
	ServiceName string
}

// Serve starts the function
func (l *LocalService) Serve() error {
	return l.Function()
}

// String returns user-friendly service name
func (l *LocalService) String() string {
	return l.ServiceName
}

// Name returns unique service name
func (l *LocalService) Name() string {
	return l.ServiceName
}

// ServiceFunc is a service function
type ServiceFunc func() error

const (
	stateCreated = iota
	stateStarted = iota
)
<commit_msg>Fixed incorrect access of logger.<commit_after>/*
Copyright 2015 Gravitational, Inc.

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 service

import (
	"context"
	"sync"

	"github.com/gravitational/teleport/lib/utils"

	"github.com/gravitational/trace"
	"github.com/sirupsen/logrus"
)

// Supervisor implements the simple service logic - registering
// service functions and de-registering the service goroutines
type Supervisor interface {
	// Register adds the service to the pool, if supervisor is in
	// the started state, the service will be started immediately
	// otherwise, it will be started after Start() has been called
	Register(srv Service)

	// RegisterFunc creates a service from function spec and registers
	// it within the system
	RegisterFunc(name string, fn ServiceFunc)

	// ServiceCount returns the number of registered and actively running
	// services
	ServiceCount() int

	// Start starts all unstarted services
	Start() error

	// Wait waits until all services exit
	Wait() error

	// Run starts and waits for the service to complete
	// it's a combinatioin Start() and Wait()
	Run() error

	// Services returns list of running services
	Services() []string

	// BroadcastEvent generates event and broadcasts it to all
	// interested parties
	BroadcastEvent(Event)

	// WaitForEvent waits for event to be broadcasted, if the event
	// was already broadcasted, payloadC will receive current event immediately
	// CLose 'cancelC' channel to force WaitForEvent to return prematurely
	WaitForEvent(name string, eventC chan Event, cancelC chan struct{})
}

type LocalSupervisor struct {
	state int
	sync.Mutex
	wg           *sync.WaitGroup
	services     []Service
	errors       []error
	events       map[string]Event
	eventsC      chan Event
	eventWaiters map[string][]*waiter
	closeContext context.Context
	signalClose  context.CancelFunc
}

// NewSupervisor returns new instance of initialized supervisor
func NewSupervisor() Supervisor {
	closeContext, cancel := context.WithCancel(context.TODO())
	srv := &LocalSupervisor{
		services:     []Service{},
		wg:           &sync.WaitGroup{},
		events:       map[string]Event{},
		eventsC:      make(chan Event, 1024),
		eventWaiters: make(map[string][]*waiter),
		closeContext: closeContext,
		signalClose:  cancel,
	}
	go srv.fanOut()
	return srv
}

// Event is a special service event that can be generated
// by various goroutines in the supervisor
type Event struct {
	Name    string
	Payload interface{}
}

func (e *Event) String() string {
	return e.Name
}

func (s *LocalSupervisor) Register(srv Service) {
	log.WithFields(logrus.Fields{"service": srv.Name()}).Debugf("Adding service to supervisor.")
	s.Lock()
	defer s.Unlock()
	s.services = append(s.services, srv)

	if s.state == stateStarted {
		s.serve(srv)
	}
}

// ServiceCount returns the number of registered and actively running services
func (s *LocalSupervisor) ServiceCount() int {
	s.Lock()
	defer s.Unlock()
	return len(s.services)
}

// RegisterFunc creates a service from function spec and registers
// it within the system
func (s *LocalSupervisor) RegisterFunc(name string, fn ServiceFunc) {
	s.Register(&LocalService{Function: fn, ServiceName: name})
}

// RemoveService removes service from supervisor tracking list
func (s *LocalSupervisor) RemoveService(srv Service) error {
	l := logrus.WithFields(logrus.Fields{"service": srv.Name()})
	s.Lock()
	defer s.Unlock()
	for i, el := range s.services {
		if el == srv {
			s.services = append(s.services[:i], s.services[i+1:]...)
			l.Debugf("Service is completed and removed.")
			return nil
		}
	}
	l.Warningf("Service is completed but not found.")
	return trace.NotFound("service %v is not found", srv)
}

func (s *LocalSupervisor) serve(srv Service) {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		defer s.RemoveService(srv)
		log.WithFields(logrus.Fields{"service": srv.Name()}).Debugf("Service has started.")
		err := srv.Serve()
		if err != nil {
			utils.FatalError(err)
		}
	}()
}

func (s *LocalSupervisor) Start() error {
	s.Lock()
	defer s.Unlock()
	s.state = stateStarted

	if len(s.services) == 0 {
		log.Warning("Supervisor has no services to run. Exiting.")
		return nil
	}

	for _, srv := range s.services {
		s.serve(srv)
	}

	return nil
}

func (s *LocalSupervisor) Services() []string {
	s.Lock()
	defer s.Unlock()

	out := make([]string, len(s.services))

	for i, srv := range s.services {
		out[i] = srv.Name()
	}
	return out
}

func (s *LocalSupervisor) Wait() error {
	defer s.signalClose()
	s.wg.Wait()
	return nil
}

func (s *LocalSupervisor) Run() error {
	if err := s.Start(); err != nil {
		return trace.Wrap(err)
	}
	return s.Wait()
}

func (s *LocalSupervisor) BroadcastEvent(event Event) {
	s.Lock()
	defer s.Unlock()
	s.events[event.Name] = event
	log.WithFields(logrus.Fields{"event": event.String()}).Debugf("Broadcasting event.")

	go func() {
		s.eventsC <- event
	}()
}

func (s *LocalSupervisor) WaitForEvent(name string, eventC chan Event, cancelC chan struct{}) {
	s.Lock()
	defer s.Unlock()

	waiter := &waiter{eventC: eventC, cancelC: cancelC}
	event, ok := s.events[name]
	if ok {
		go s.notifyWaiter(waiter, event)
		return
	}
	s.eventWaiters[name] = append(s.eventWaiters[name], waiter)
}

func (s *LocalSupervisor) getWaiters(name string) []*waiter {
	s.Lock()
	defer s.Unlock()

	waiters := s.eventWaiters[name]
	out := make([]*waiter, len(waiters))
	for i := range waiters {
		out[i] = waiters[i]
	}
	return out
}

func (s *LocalSupervisor) notifyWaiter(w *waiter, event Event) {
	select {
	case w.eventC <- event:
	case <-w.cancelC:
	}
}

func (s *LocalSupervisor) fanOut() {
	for {
		select {
		case event := <-s.eventsC:
			waiters := s.getWaiters(event.Name)
			for _, waiter := range waiters {
				go s.notifyWaiter(waiter, event)
			}
		case <-s.closeContext.Done():
			return
		}
	}
}

type waiter struct {
	eventC  chan Event
	cancelC chan struct{}
}

// Service is a running teleport service function
type Service interface {
	// Serve starts the function
	Serve() error
	// String returns user-friendly description of service
	String() string
	// Name returns service name
	Name() string
}

// LocalService is a locally defined service
type LocalService struct {
	// Function is a function to call
	Function ServiceFunc
	// ServiceName is a service name
	ServiceName string
}

// Serve starts the function
func (l *LocalService) Serve() error {
	return l.Function()
}

// String returns user-friendly service name
func (l *LocalService) String() string {
	return l.ServiceName
}

// Name returns unique service name
func (l *LocalService) Name() string {
	return l.ServiceName
}

// ServiceFunc is a service function
type ServiceFunc func() error

const (
	stateCreated = iota
	stateStarted = iota
)
<|endoftext|>


<commit_before>package hcsshim

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/Microsoft/go-winio"
)

var errorIterationCanceled = errors.New("")

func openFileOrDir(path string, mode uint32, createDisposition uint32) (file *os.File, err error) {
	winPath, err := syscall.UTF16FromString(path)
	if err != nil {
		return
	}
	h, err := syscall.CreateFile(&winPath[0], mode, syscall.FILE_SHARE_READ, nil, createDisposition, syscall.FILE_FLAG_BACKUP_SEMANTICS, 0)
	if err != nil {
		err = &os.PathError{"open", path, err}
		return
	}
	file = os.NewFile(uintptr(h), path)
	return
}

type fileEntry struct {
	path string
	fi   os.FileInfo
	err  error
}

type LegacyLayerReader struct {
	root         string
	result       chan *fileEntry
	proceed      chan bool
	currentFile  *os.File
	backupReader *winio.BackupFileReader
	isTP4Format  bool
}

// NewLegacyLayerReader returns a new LayerReader that can read the Windows
// TP4 transport format from disk.
func NewLegacyLayerReader(root string) *LegacyLayerReader {
	r := &LegacyLayerReader{
		root:        root,
		result:      make(chan *fileEntry),
		proceed:     make(chan bool),
		isTP4Format: IsTP4(),
	}
	go r.walk()
	return r
}

func readTombstones(path string) (map[string]([]string), error) {
	tf, err := os.Open(filepath.Join(path, "tombstones.txt"))
	if err != nil {
		return nil, err
	}
	defer tf.Close()
	s := bufio.NewScanner(tf)
	if !s.Scan() || s.Text() != "\xef\xbb\xbfVersion 1.0" {
		return nil, errors.New("Invalid tombstones file")
	}

	ts := make(map[string]([]string))
	for s.Scan() {
		t := s.Text()[1:] // skip leading `\`
		dir := filepath.Dir(t)
		ts[dir] = append(ts[dir], t)
	}
	if err = s.Err(); err != nil {
		return nil, err
	}

	return ts, nil
}

func (r *LegacyLayerReader) walk() {
	defer close(r.result)
	if !<-r.proceed {
		return
	}

	ts, err := readTombstones(r.root)
	if err != nil {
		goto ErrorLoop
	}

	err = filepath.Walk(r.root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if path == r.root || path == filepath.Join(r.root, "tombstones.txt") || strings.HasSuffix(path, ".$wcidirs$") {
			return nil
		}
		r.result <- &fileEntry{path, info, nil}
		if !<-r.proceed {
			return errorIterationCanceled
		}

		// List all the tombstones.
		if info.IsDir() {
			relPath, err := filepath.Rel(r.root, path)
			if err != nil {
				return err
			}
			if dts, ok := ts[relPath]; ok {
				for _, t := range dts {
					r.result <- &fileEntry{t, nil, nil}
					if !<-r.proceed {
						return errorIterationCanceled
					}
				}
			}
		}
		return nil
	})
	if err == errorIterationCanceled {
		return
	}
	if err == nil {
		err = io.EOF
	}

ErrorLoop:
	for {
		r.result <- &fileEntry{err: err}
		if !<-r.proceed {
			break
		}
	}
}

func (r *LegacyLayerReader) reset() {
	if r.backupReader != nil {
		r.backupReader.Close()
		r.backupReader = nil
	}
	if r.currentFile != nil {
		r.currentFile.Close()
		r.currentFile = nil
	}
}

func findBackupStreamSize(r io.Reader) (int64, error) {
	br := winio.NewBackupStreamReader(r)
	for {
		hdr, err := br.Next()
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return 0, err
		}
		if hdr.Id == winio.BackupData {
			return hdr.Size, nil
		}
	}
}

func (r *LegacyLayerReader) Next() (path string, size int64, fileInfo *winio.FileBasicInfo, err error) {
	r.reset()
	r.proceed <- true
	fe := <-r.result
	if fe == nil {
		err = errors.New("LegacyLayerReader closed")
		return
	}
	if fe.err != nil {
		err = fe.err
		return
	}

	path, err = filepath.Rel(r.root, fe.path)
	if err != nil {
		return
	}

	if fe.fi == nil {
		// This is a tombstone. Return a nil fileInfo.
		return
	}

	if fe.fi.IsDir() && strings.HasPrefix(path, `Files\`) {
		fe.path += ".$wcidirs$"
	}

	f, err := openFileOrDir(fe.path, syscall.GENERIC_READ, syscall.OPEN_EXISTING)
	if err != nil {
		return
	}
	defer func() {
		if f != nil {
			f.Close()
		}
	}()

	fileInfo, err = winio.GetFileBasicInfo(f)
	if err != nil {
		return
	}

	if !strings.HasPrefix(path, `Files\`) {
		size = fe.fi.Size()
		r.backupReader = winio.NewBackupFileReader(f, false)
		if path == "Hives" || path == "Files" {
			// The Hives directory has a non-deterministic file time because of the
			// nature of the import process. Use the times from System_Delta.
			var g *os.File
			g, err = os.Open(filepath.Join(r.root, `Hives\System_Delta`))
			if err != nil {
				return
			}
			attr := fileInfo.FileAttributes
			fileInfo, err = winio.GetFileBasicInfo(g)
			g.Close()
			if err != nil {
				return
			}
			fileInfo.FileAttributes = attr
		}

		// The creation time and access time get reset for files outside of the Files path.
		fileInfo.CreationTime = fileInfo.LastWriteTime
		fileInfo.LastAccessTime = fileInfo.LastWriteTime

	} else {
		beginning := int64(0)
		if !r.isTP4Format {
			// In TP5, the file attributes were added before the backup stream
			var attr uint32
			err = binary.Read(f, binary.LittleEndian, &attr)
			if err != nil {
				return
			}
			fileInfo.FileAttributes = uintptr(attr)
			beginning = 4
		}

		// Find the accurate file size.
		if !fe.fi.IsDir() {
			size, err = findBackupStreamSize(f)
			if err != nil {
				err = &os.PathError{"findBackupStreamSize", fe.path, err}
				return
			}
		}

		// Return back to the beginning of the backup stream.
		_, err = f.Seek(beginning, 0)
		if err != nil {
			return
		}
	}

	r.currentFile = f
	f = nil
	return
}

func (r *LegacyLayerReader) Read(b []byte) (int, error) {
	if r.backupReader == nil {
		if r.currentFile == nil {
			return 0, io.EOF
		}
		return r.currentFile.Read(b)
	}
	return r.backupReader.Read(b)
}

func (r *LegacyLayerReader) Close() error {
	r.proceed <- false
	<-r.result
	r.reset()
	return nil
}

type LegacyLayerWriter struct {
	root         string
	currentFile  *os.File
	backupWriter *winio.BackupFileWriter
	tombstones   []string
	isTP4Format  bool
}

// NewLegacyLayerWriter returns a LayerWriter that can write the TP4 transport format
// to disk.
func NewLegacyLayerWriter(root string) *LegacyLayerWriter {
	return &LegacyLayerWriter{
		root:        root,
		isTP4Format: IsTP4(),
	}
}

func (w *LegacyLayerWriter) reset() {
	if w.backupWriter != nil {
		w.backupWriter.Close()
		w.backupWriter = nil
	}
	if w.currentFile != nil {
		w.currentFile.Close()
		w.currentFile = nil
	}
}

func (w *LegacyLayerWriter) Add(name string, fileInfo *winio.FileBasicInfo) error {
	w.reset()
	path := filepath.Join(w.root, name)

	createDisposition := uint32(syscall.CREATE_NEW)
	if (fileInfo.FileAttributes & syscall.FILE_ATTRIBUTE_DIRECTORY) != 0 {
		err := os.Mkdir(path, 0)
		if err != nil {
			return err
		}
		if strings.HasPrefix(name, `Files\`) {
			path += ".$wcidirs$"
		} else {
			createDisposition = syscall.OPEN_EXISTING
		}
	}

	f, err := openFileOrDir(path, syscall.GENERIC_READ|syscall.GENERIC_WRITE, createDisposition)
	if err != nil {
		return err
	}
	defer func() {
		if f != nil {
			f.Close()
			os.Remove(path)
		}
	}()

	strippedFi := *fileInfo
	strippedFi.FileAttributes = 0
	err = winio.SetFileBasicInfo(f, &strippedFi)
	if err != nil {
		return err
	}

	if !strings.HasPrefix(name, `Files\`) {
		w.backupWriter = winio.NewBackupFileWriter(f, false)
	} else {
		if !w.isTP4Format {
			// In TP5, the file attributes were added to the header
			err = binary.Write(f, binary.LittleEndian, uint32(fileInfo.FileAttributes))
			if err != nil {
				return err
			}
		}
	}

	w.currentFile = f
	f = nil
	return nil
}

func (w *LegacyLayerWriter) Remove(name string) error {
	w.tombstones = append(w.tombstones, name)
	return nil
}

func (w *LegacyLayerWriter) Write(b []byte) (int, error) {
	if w.backupWriter == nil {
		if w.currentFile == nil {
			return 0, errors.New("closed")
		}
		return w.currentFile.Write(b)
	}
	return w.backupWriter.Write(b)
}

func (w *LegacyLayerWriter) Close() error {
	w.reset()
	tf, err := os.Create(filepath.Join(w.root, "tombstones.txt"))
	if err != nil {
		return err
	}
	defer tf.Close()
	_, err = tf.Write([]byte("\xef\xbb\xbfVersion 1.0\n"))
	if err != nil {
		return err
	}
	for _, t := range w.tombstones {
		_, err = tf.Write([]byte(filepath.Join(`\`, t) + "\n"))
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Fix legacy export<commit_after>package hcsshim

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/Microsoft/go-winio"
)

var errorIterationCanceled = errors.New("")

func openFileOrDir(path string, mode uint32, createDisposition uint32) (file *os.File, err error) {
	winPath, err := syscall.UTF16FromString(path)
	if err != nil {
		return
	}
	h, err := syscall.CreateFile(&winPath[0], mode, syscall.FILE_SHARE_READ, nil, createDisposition, syscall.FILE_FLAG_BACKUP_SEMANTICS, 0)
	if err != nil {
		err = &os.PathError{"open", path, err}
		return
	}
	file = os.NewFile(uintptr(h), path)
	return
}

type fileEntry struct {
	path string
	fi   os.FileInfo
	err  error
}

type LegacyLayerReader struct {
	root         string
	result       chan *fileEntry
	proceed      chan bool
	currentFile  *os.File
	backupReader *winio.BackupFileReader
	isTP4Format  bool
}

// NewLegacyLayerReader returns a new LayerReader that can read the Windows
// TP4 transport format from disk.
func NewLegacyLayerReader(root string) *LegacyLayerReader {
	r := &LegacyLayerReader{
		root:        root,
		result:      make(chan *fileEntry),
		proceed:     make(chan bool),
		isTP4Format: IsTP4(),
	}
	go r.walk()
	return r
}

func readTombstones(path string) (map[string]([]string), error) {
	tf, err := os.Open(filepath.Join(path, "tombstones.txt"))
	if err != nil {
		return nil, err
	}
	defer tf.Close()
	s := bufio.NewScanner(tf)
	if !s.Scan() || s.Text() != "\xef\xbb\xbfVersion 1.0" {
		return nil, errors.New("Invalid tombstones file")
	}

	ts := make(map[string]([]string))
	for s.Scan() {
		t := s.Text()[1:] // skip leading `\`
		dir := filepath.Dir(t)
		ts[dir] = append(ts[dir], t)
	}
	if err = s.Err(); err != nil {
		return nil, err
	}

	return ts, nil
}

func (r *LegacyLayerReader) walk() {
	defer close(r.result)
	if !<-r.proceed {
		return
	}

	ts, err := readTombstones(r.root)
	if err != nil {
		goto ErrorLoop
	}

	err = filepath.Walk(r.root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if path == r.root || path == filepath.Join(r.root, "tombstones.txt") || strings.HasSuffix(path, ".$wcidirs$") {
			return nil
		}
		r.result <- &fileEntry{path, info, nil}
		if !<-r.proceed {
			return errorIterationCanceled
		}

		// List all the tombstones.
		if info.IsDir() {
			relPath, err := filepath.Rel(r.root, path)
			if err != nil {
				return err
			}
			if dts, ok := ts[relPath]; ok {
				for _, t := range dts {
					r.result <- &fileEntry{t, nil, nil}
					if !<-r.proceed {
						return errorIterationCanceled
					}
				}
			}
		}
		return nil
	})
	if err == errorIterationCanceled {
		return
	}
	if err == nil {
		err = io.EOF
	}

ErrorLoop:
	for {
		r.result <- &fileEntry{err: err}
		if !<-r.proceed {
			break
		}
	}
}

func (r *LegacyLayerReader) reset() {
	if r.backupReader != nil {
		r.backupReader.Close()
		r.backupReader = nil
	}
	if r.currentFile != nil {
		r.currentFile.Close()
		r.currentFile = nil
	}
}

func findBackupStreamSize(r io.Reader) (int64, error) {
	br := winio.NewBackupStreamReader(r)
	for {
		hdr, err := br.Next()
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return 0, err
		}
		if hdr.Id == winio.BackupData {
			return hdr.Size, nil
		}
	}
}

func (r *LegacyLayerReader) Next() (path string, size int64, fileInfo *winio.FileBasicInfo, err error) {
	r.reset()
	r.proceed <- true
	fe := <-r.result
	if fe == nil {
		err = errors.New("LegacyLayerReader closed")
		return
	}
	if fe.err != nil {
		err = fe.err
		return
	}

	path, err = filepath.Rel(r.root, fe.path)
	if err != nil {
		return
	}

	if fe.fi == nil {
		// This is a tombstone. Return a nil fileInfo.
		return
	}

	if fe.fi.IsDir() && strings.HasPrefix(path, `Files\`) {
		fe.path += ".$wcidirs$"
	}

	f, err := openFileOrDir(fe.path, syscall.GENERIC_READ, syscall.OPEN_EXISTING)
	if err != nil {
		return
	}
	defer func() {
		if f != nil {
			f.Close()
		}
	}()

	fileInfo, err = winio.GetFileBasicInfo(f)
	if err != nil {
		return
	}

	if !strings.HasPrefix(path, `Files\`) {
		size = fe.fi.Size()
		r.backupReader = winio.NewBackupFileReader(f, false)
		if path == "Hives" || path == "Files" {
			// The Hives directory has a non-deterministic file time because of the
			// nature of the import process. Use the times from System_Delta.
			var g *os.File
			g, err = os.Open(filepath.Join(r.root, `Hives\System_Delta`))
			if err != nil {
				return
			}
			attr := fileInfo.FileAttributes
			fileInfo, err = winio.GetFileBasicInfo(g)
			g.Close()
			if err != nil {
				return
			}
			fileInfo.FileAttributes = attr
		}

		// The creation time and access time get reset for files outside of the Files path.
		fileInfo.CreationTime = fileInfo.LastWriteTime
		fileInfo.LastAccessTime = fileInfo.LastWriteTime

	} else {
		beginning := int64(0)
		if !r.isTP4Format {
			// In TP5, the file attributes were added before the backup stream
			var attr uint32
			err = binary.Read(f, binary.LittleEndian, &attr)
			if err != nil {
				return
			}
			fileInfo.FileAttributes = uintptr(attr)
			beginning = 4
		}

		// Find the accurate file size.
		if !fe.fi.IsDir() {
			size, err = findBackupStreamSize(f)
			if err != nil {
				err = &os.PathError{"findBackupStreamSize", fe.path, err}
				return
			}
		}

		// Return back to the beginning of the backup stream.
		_, err = f.Seek(beginning, 0)
		if err != nil {
			return
		}
	}

	r.currentFile = f
	f = nil
	return
}

func (r *LegacyLayerReader) Read(b []byte) (int, error) {
	if r.backupReader == nil {
		if r.currentFile == nil {
			return 0, io.EOF
		}
		return r.currentFile.Read(b)
	}
	return r.backupReader.Read(b)
}

func (r *LegacyLayerReader) Close() error {
	r.proceed <- false
	<-r.result
	r.reset()
	return nil
}

type LegacyLayerWriter struct {
	root         string
	currentFile  *os.File
	backupWriter *winio.BackupFileWriter
	tombstones   []string
	isTP4Format  bool
}

// NewLegacyLayerWriter returns a LayerWriter that can write the TP4 transport format
// to disk.
func NewLegacyLayerWriter(root string) *LegacyLayerWriter {
	return &LegacyLayerWriter{
		root:        root,
		isTP4Format: IsTP4(),
	}
}

func (w *LegacyLayerWriter) reset() {
	if w.backupWriter != nil {
		w.backupWriter.Close()
		w.backupWriter = nil
	}
	if w.currentFile != nil {
		w.currentFile.Close()
		w.currentFile = nil
	}
}

func (w *LegacyLayerWriter) Add(name string, fileInfo *winio.FileBasicInfo) error {
	w.reset()
	path := filepath.Join(w.root, name)

	createDisposition := uint32(syscall.CREATE_NEW)
	if (fileInfo.FileAttributes & syscall.FILE_ATTRIBUTE_DIRECTORY) != 0 {
		err := os.Mkdir(path, 0)
		if err != nil {
			return err
		}
		path += ".$wcidirs$"
	}

	f, err := openFileOrDir(path, syscall.GENERIC_READ|syscall.GENERIC_WRITE, createDisposition)
	if err != nil {
		return err
	}
	defer func() {
		if f != nil {
			f.Close()
			os.Remove(path)
		}
	}()

	strippedFi := *fileInfo
	strippedFi.FileAttributes = 0
	err = winio.SetFileBasicInfo(f, &strippedFi)
	if err != nil {
		return err
	}

	if strings.HasPrefix(name, `Hives\`) {
		w.backupWriter = winio.NewBackupFileWriter(f, false)
	} else {
		if !w.isTP4Format {
			// In TP5, the file attributes were added to the header
			err = binary.Write(f, binary.LittleEndian, uint32(fileInfo.FileAttributes))
			if err != nil {
				return err
			}
		}
	}

	w.currentFile = f
	f = nil
	return nil
}

func (w *LegacyLayerWriter) Remove(name string) error {
	w.tombstones = append(w.tombstones, name)
	return nil
}

func (w *LegacyLayerWriter) Write(b []byte) (int, error) {
	if w.backupWriter == nil {
		if w.currentFile == nil {
			return 0, errors.New("closed")
		}
		return w.currentFile.Write(b)
	}
	return w.backupWriter.Write(b)
}

func (w *LegacyLayerWriter) Close() error {
	w.reset()
	tf, err := os.Create(filepath.Join(w.root, "tombstones.txt"))
	if err != nil {
		return err
	}
	defer tf.Close()
	_, err = tf.Write([]byte("\xef\xbb\xbfVersion 1.0\n"))
	if err != nil {
		return err
	}
	for _, t := range w.tombstones {
		_, err = tf.Write([]byte(filepath.Join(`\`, t) + "\n"))
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sync"

	"github.com/nsf/termbox-go"

	"github.com/prattmic/lesser/lineio"
	"github.com/prattmic/lesser/sortedmap"
)

type size struct {
	x int
	y int
}

type Event int

const (
	// EventQuit requests an application exit.
	EventQuit Event = iota

	// EventRefresh requests a display refresh.
	EventRefresh
)

type Mode int

const (
	// ModeNormal is the standard mode, allowing file navigation.
	ModeNormal Mode = iota

	// ModeSearchEntry is search entry mode. Key presses are added
	// to the search string.
	ModeSearchEntry
)

type Lesser struct {
	// src is the source file being displayed.
	src lineio.LineReader

	// tabStop is the number of spaces per tab.
	tabStop int

	// events is used to notify the main goroutine of events.
	events chan Event

	// mu locks the fields below.
	mu sync.Mutex

	// size is the size of the file display.
	// There is a statusbar beneath the display.
	size size

	// line is the line number of the first line of the display.
	line int64

	// mode is the viewer mode.
	mode Mode

	// regexp is the search regexp specified by the user.
	// Must only be modified by the event goroutine.
	regexp string

	// searchResults are the results for the current search.
	// They should be highlighted.
	searchResults searchResults
}

// lastLine returns the last line on the display.  It may be beyond the end
// of the file, if the file is short enough.
// mu must be held on call.
func (l *Lesser) lastLine() int64 {
	return l.line + int64(l.size.y) - 1
}

// Scroll describes a scroll action.
type Scroll int

const (
	// ScrollTop goes to the first line.
	ScrollTop Scroll = iota
	// ScrollBottom goes to the last line.
	ScrollBottom
	// ScrollUp goes up one line.
	ScrollUp
	// ScrollDown goes down one line.
	ScrollDown
	// ScrollUpPage goes up one page full.
	ScrollUpPage
	// ScrollDownPage goes down one page full.
	ScrollDownPage
	// ScrollUpHalfPage goes up one half page full.
	ScrollUpHalfPage
	// ScrollDownHalfPage goes down one half page full.
	ScrollDownHalfPage
)

// scrollLine tries to scroll the display to the given line,
// but will not scroll beyond the first or last lines in the file.
// l.mu must be held when calling scrollLine.
func (l *Lesser) scrollLine(dest int64) {
	var delta int64
	if dest > l.line {
		delta = 1
	} else {
		delta = -1
	}

	for l.line != dest && l.line+delta > 0 && l.src.LineExists(l.lastLine()+delta) {
		l.line += delta
	}
}

// scroll moves the display based on the passed scroll action, without
// going past the beginning or end of the file.
func (l *Lesser) scroll(s Scroll) {
	l.mu.Lock()
	defer l.mu.Unlock()

	var dest int64
	switch s {
	case ScrollTop:
		dest = 1
	case ScrollBottom:
		// Just try to go to int64 max.
		dest = 0x7fffffffffffffff
	case ScrollUp:
		dest = l.line - 1
	case ScrollDown:
		dest = l.line + 1
	case ScrollUpPage:
		dest = l.line - int64(l.size.y)
	case ScrollDownPage:
		dest = l.line + int64(l.size.y)
	case ScrollUpHalfPage:
		dest = l.line - int64(l.size.y)/2
	case ScrollDownHalfPage:
		dest = l.line + int64(l.size.y)/2
	}

	l.scrollLine(dest)
}

func (l *Lesser) handleEvent(e termbox.Event) {
	l.mu.Lock()
	mode := l.mode
	l.mu.Unlock()

	if e.Type != termbox.EventKey {
		return
	}

	c := e.Ch
	k := e.Key
	// Key is only valid is Ch is 0
	if c != 0 {
		k = 0
	}

	switch mode {
	case ModeNormal:
		switch {
		case c == 'q':
			l.events <- EventQuit
		case c == 'j':
			l.scroll(ScrollDown)
			l.events <- EventRefresh
		case c == 'k':
			l.scroll(ScrollUp)
			l.events <- EventRefresh
		case c == 'g':
			l.scroll(ScrollTop)
			l.events <- EventRefresh
		case c == 'G':
			l.scroll(ScrollBottom)
			l.events <- EventRefresh
		case k == termbox.KeyPgup:
			l.scroll(ScrollUpPage)
			l.events <- EventRefresh
		case k == termbox.KeyPgdn:
			l.scroll(ScrollDownPage)
			l.events <- EventRefresh
		case k == termbox.KeyCtrlU:
			l.scroll(ScrollUpHalfPage)
			l.events <- EventRefresh
		case k == termbox.KeyCtrlD:
			l.scroll(ScrollDownHalfPage)
			l.events <- EventRefresh
		case c == '/':
			l.mu.Lock()
			l.mode = ModeSearchEntry
			l.mu.Unlock()
			l.events <- EventRefresh
		case c == 'n':
			l.mu.Lock()
			if r, ok := l.searchResults.Next(l.line); ok {
				l.scrollLine(r.line)
				l.events <- EventRefresh
			}
			l.mu.Unlock()
		case c == 'N':
			l.mu.Lock()
			if r, ok := l.searchResults.Prev(l.line); ok {
				l.scrollLine(r.line)
				l.events <- EventRefresh
			}
			l.mu.Unlock()
		}
	case ModeSearchEntry:
		switch {
		case k == termbox.KeyEnter:
			r := l.search(l.regexp)
			l.mu.Lock()
			l.mode = ModeNormal
			l.regexp = ""
			l.searchResults = r
			l.mu.Unlock()
			l.events <- EventRefresh
		default:
			l.mu.Lock()
			l.regexp += string(c)
			l.mu.Unlock()
			l.events <- EventRefresh
		}
	}
}

func (l *Lesser) listenEvents() {
	for {
		e := termbox.PollEvent()
		l.handleEvent(e)
	}
}

// searchResult describes search matches on a single line.
type searchResult struct {
	line    int64
	matches [][]int
	err     error
}

// matchesChar returns true if the search result contains a match for
// character index c.
func (s searchResult) matchesChar(c int) bool {
	for _, match := range s.matches {
		if len(match) < 2 {
			continue
		}

		if c >= match[0] && c < match[1] {
			return true
		}
	}
	return false
}

type searchResults struct {
	// mu locks the fields below.
	mu sync.Mutex

	// lines maps search results for a specific line to an index in results.
	lines sortedmap.Map

	// results contains the actual search results, in no particular order.
	results []searchResult
}

func NewSearchResults() searchResults {
	return searchResults{
		lines: sortedmap.NewMap(),
	}
}

// Add adds a result.
func (s *searchResults) Add(r searchResult) {
	s.mu.Lock()
	defer s.mu.Unlock()

	i := int64(len(s.results))
	s.results = append(s.results, r)
	s.lines.Insert(r.line, i)
}

// Get finds the result for a specific line, returning ok if found
func (s *searchResults) Get(line int64) (searchResult, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if i, ok := s.lines.Get(line); ok {
		return s.results[i], true
	}

	return searchResult{}, false
}

// Next returns the search result for the nearest line after line,
// noninclusive, if one exists.
func (s *searchResults) Next(line int64) (searchResult, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, i, err := s.lines.NearestGreater(line)
	if err != nil {
		// Probably ErrNoSuchKey, aka none found.
		return searchResult{}, false
	}

	return s.results[i], true
}

// Prev returns the search result for the nearest line before line,
// noninclusive, if one exists.
func (s *searchResults) Prev(line int64) (searchResult, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Search for line - 1, since it may be equal.
	_, i, err := s.lines.NearestLessEqual(line - 1)
	if err != nil {
		// Probably ErrNoSuchKey, aka none found.
		return searchResult{}, false
	}

	return s.results[i], true
}

func (l *Lesser) search(s string) searchResults {
	reg, err := regexp.Compile(s)
	if err != nil {
		// TODO(prattmic): display a better error
		log.Printf("regexp failed to compile: %v", err)
		return NewSearchResults()
	}

	resultChan := make(chan searchResult, 100)

	searchLine := func(line int64) {
		r, err := l.src.SearchLine(reg, line)
		if err != nil {
			r = nil
		}

		resultChan <- searchResult{
			line:    line,
			matches: r,
			err:     err,
		}
	}

	nextLine := int64(1)
	// Spawn initial search goroutines
	for ; nextLine <= 5; nextLine++ {
		go searchLine(nextLine)
	}

	results := NewSearchResults()

	var count int64

	waitResult := func() searchResult {
		ret := <-resultChan
		count += 1

		// Only store results with matches.
		if len(ret.matches) > 0 {
			results.Add(ret)
		}

		return ret
	}

	// Collect results, start searching next lines until we start
	// hitting EOF.
	for {
		r := waitResult()

		// We started hitting errors on a previous line,
		// there is no reason to search later lines.
		if r.err != nil {
			break
		}

		go searchLine(nextLine)
		nextLine++
	}

	// Collect the remaing results.
	for count < nextLine-1 {
		waitResult()
	}

	return results
}

// statusBar renders the status bar.
// mu must be held on call.
func (l *Lesser) statusBar() {
	// The statusbar is just below the display.

	// Clear the statusbar
	for i := 0; i < l.size.x; i++ {
		termbox.SetCell(i, l.size.y, ' ', 0, 0)
	}

	switch l.mode {
	case ModeNormal:
		// Just a colon and a cursor
		termbox.SetCell(0, l.size.y, ':', 0, 0)
		termbox.SetCursor(1, l.size.y)
	case ModeSearchEntry:
		// / and search string
		termbox.SetCell(0, l.size.y, '/', 0, 0)
		for i, c := range l.regexp {
			termbox.SetCell(1+i, l.size.y, c, 0, 0)
		}
		termbox.SetCursor(1+len(l.regexp), l.size.y)
	}
}

// alignUp aligns n up to the next multiple of divisor.
func alignUp(n, divisor int) int {
	return n + (divisor - (n % divisor))
}

func (l *Lesser) refreshScreen() error {
	l.mu.Lock()
	defer l.mu.Unlock()

	for y := 0; y < l.size.y; y++ {
		buf := make([]byte, l.size.x)
		line := l.line + int64(y)

		_, err := l.src.ReadLine(buf, line)
		// EOF just means the line was shorter than the display.
		if err != nil && err != io.EOF {
			return err
		}

		highlight, ok := l.searchResults.Get(line)

		var displayColumn int
		for i, c := range buf {
			if displayColumn >= l.size.x {
				break
			}

			fg := termbox.ColorDefault
			bg := termbox.ColorDefault

			// Highlight matches
			if ok && highlight.matchesChar(i) {
				fg = termbox.ColorBlack
				bg = termbox.ColorWhite
			}

			if c == '\t' {
				// Tabs align the display up to the next
				// multiple of tabstop.
				next := alignUp(displayColumn, l.tabStop)

				// Clear the tab spaces
				for j := displayColumn; j < next; j++ {
					termbox.SetCell(j, y, ' ', 0, 0)
				}

				displayColumn = next
			} else {
				termbox.SetCell(displayColumn, y, rune(c), fg, bg)
				displayColumn += 1
			}
		}
	}

	l.statusBar()

	termbox.Flush()

	return nil
}

func (l *Lesser) Run() {
	go l.listenEvents()

	err := l.refreshScreen()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to refresh screen: %v\n", err)
		return
	}

	for {
		e := <-l.events

		switch e {
		case EventQuit:
			return
		case EventRefresh:
			err = l.refreshScreen()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Failed to refresh screen: %v\n", err)
				return
			}
		}
	}
}

func NewLesser(f *os.File, ts int) Lesser {
	x, y := termbox.Size()

	return Lesser{
		src:     lineio.NewLineReader(f),
		tabStop: ts,
		// Save one line for statusbar.
		size:   size{x: x, y: y - 1},
		line:   1,
		events: make(chan Event, 1),
		mode:   ModeNormal,
	}
}
<commit_msg>Jump to nearest result when search completes<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sync"

	"github.com/nsf/termbox-go"

	"github.com/prattmic/lesser/lineio"
	"github.com/prattmic/lesser/sortedmap"
)

type size struct {
	x int
	y int
}

type Event int

const (
	// EventQuit requests an application exit.
	EventQuit Event = iota

	// EventRefresh requests a display refresh.
	EventRefresh
)

type Mode int

const (
	// ModeNormal is the standard mode, allowing file navigation.
	ModeNormal Mode = iota

	// ModeSearchEntry is search entry mode. Key presses are added
	// to the search string.
	ModeSearchEntry
)

type Lesser struct {
	// src is the source file being displayed.
	src lineio.LineReader

	// tabStop is the number of spaces per tab.
	tabStop int

	// events is used to notify the main goroutine of events.
	events chan Event

	// mu locks the fields below.
	mu sync.Mutex

	// size is the size of the file display.
	// There is a statusbar beneath the display.
	size size

	// line is the line number of the first line of the display.
	line int64

	// mode is the viewer mode.
	mode Mode

	// regexp is the search regexp specified by the user.
	// Must only be modified by the event goroutine.
	regexp string

	// searchResults are the results for the current search.
	// They should be highlighted.
	searchResults searchResults
}

// lastLine returns the last line on the display.  It may be beyond the end
// of the file, if the file is short enough.
// mu must be held on call.
func (l *Lesser) lastLine() int64 {
	return l.line + int64(l.size.y) - 1
}

// Scroll describes a scroll action.
type Scroll int

const (
	// ScrollTop goes to the first line.
	ScrollTop Scroll = iota
	// ScrollBottom goes to the last line.
	ScrollBottom
	// ScrollUp goes up one line.
	ScrollUp
	// ScrollDown goes down one line.
	ScrollDown
	// ScrollUpPage goes up one page full.
	ScrollUpPage
	// ScrollDownPage goes down one page full.
	ScrollDownPage
	// ScrollUpHalfPage goes up one half page full.
	ScrollUpHalfPage
	// ScrollDownHalfPage goes down one half page full.
	ScrollDownHalfPage
)

// scrollLine tries to scroll the display to the given line,
// but will not scroll beyond the first or last lines in the file.
// l.mu must be held when calling scrollLine.
func (l *Lesser) scrollLine(dest int64) {
	var delta int64
	if dest > l.line {
		delta = 1
	} else {
		delta = -1
	}

	for l.line != dest && l.line+delta > 0 && l.src.LineExists(l.lastLine()+delta) {
		l.line += delta
	}
}

// scroll moves the display based on the passed scroll action, without
// going past the beginning or end of the file.
func (l *Lesser) scroll(s Scroll) {
	l.mu.Lock()
	defer l.mu.Unlock()

	var dest int64
	switch s {
	case ScrollTop:
		dest = 1
	case ScrollBottom:
		// Just try to go to int64 max.
		dest = 0x7fffffffffffffff
	case ScrollUp:
		dest = l.line - 1
	case ScrollDown:
		dest = l.line + 1
	case ScrollUpPage:
		dest = l.line - int64(l.size.y)
	case ScrollDownPage:
		dest = l.line + int64(l.size.y)
	case ScrollUpHalfPage:
		dest = l.line - int64(l.size.y)/2
	case ScrollDownHalfPage:
		dest = l.line + int64(l.size.y)/2
	}

	l.scrollLine(dest)
}

func (l *Lesser) handleEvent(e termbox.Event) {
	l.mu.Lock()
	mode := l.mode
	l.mu.Unlock()

	if e.Type != termbox.EventKey {
		return
	}

	c := e.Ch
	k := e.Key
	// Key is only valid is Ch is 0
	if c != 0 {
		k = 0
	}

	switch mode {
	case ModeNormal:
		switch {
		case c == 'q':
			l.events <- EventQuit
		case c == 'j':
			l.scroll(ScrollDown)
			l.events <- EventRefresh
		case c == 'k':
			l.scroll(ScrollUp)
			l.events <- EventRefresh
		case c == 'g':
			l.scroll(ScrollTop)
			l.events <- EventRefresh
		case c == 'G':
			l.scroll(ScrollBottom)
			l.events <- EventRefresh
		case k == termbox.KeyPgup:
			l.scroll(ScrollUpPage)
			l.events <- EventRefresh
		case k == termbox.KeyPgdn:
			l.scroll(ScrollDownPage)
			l.events <- EventRefresh
		case k == termbox.KeyCtrlU:
			l.scroll(ScrollUpHalfPage)
			l.events <- EventRefresh
		case k == termbox.KeyCtrlD:
			l.scroll(ScrollDownHalfPage)
			l.events <- EventRefresh
		case c == '/':
			l.mu.Lock()
			l.mode = ModeSearchEntry
			l.mu.Unlock()
			l.events <- EventRefresh
		case c == 'n':
			l.mu.Lock()
			if r, ok := l.searchResults.Next(l.line); ok {
				l.scrollLine(r.line)
				l.events <- EventRefresh
			}
			l.mu.Unlock()
		case c == 'N':
			l.mu.Lock()
			if r, ok := l.searchResults.Prev(l.line); ok {
				l.scrollLine(r.line)
				l.events <- EventRefresh
			}
			l.mu.Unlock()
		}
	case ModeSearchEntry:
		switch {
		case k == termbox.KeyEnter:
			r := l.search(l.regexp)
			l.mu.Lock()
			l.mode = ModeNormal
			l.regexp = ""
			l.searchResults = r
			// Jump to nearest result
			if r, ok := l.searchResults.Next(l.line); ok {
				l.scrollLine(r.line)
			}
			l.mu.Unlock()
			l.events <- EventRefresh
		default:
			l.mu.Lock()
			l.regexp += string(c)
			l.mu.Unlock()
			l.events <- EventRefresh
		}
	}
}

func (l *Lesser) listenEvents() {
	for {
		e := termbox.PollEvent()
		l.handleEvent(e)
	}
}

// searchResult describes search matches on a single line.
type searchResult struct {
	line    int64
	matches [][]int
	err     error
}

// matchesChar returns true if the search result contains a match for
// character index c.
func (s searchResult) matchesChar(c int) bool {
	for _, match := range s.matches {
		if len(match) < 2 {
			continue
		}

		if c >= match[0] && c < match[1] {
			return true
		}
	}
	return false
}

type searchResults struct {
	// mu locks the fields below.
	mu sync.Mutex

	// lines maps search results for a specific line to an index in results.
	lines sortedmap.Map

	// results contains the actual search results, in no particular order.
	results []searchResult
}

func NewSearchResults() searchResults {
	return searchResults{
		lines: sortedmap.NewMap(),
	}
}

// Add adds a result.
func (s *searchResults) Add(r searchResult) {
	s.mu.Lock()
	defer s.mu.Unlock()

	i := int64(len(s.results))
	s.results = append(s.results, r)
	s.lines.Insert(r.line, i)
}

// Get finds the result for a specific line, returning ok if found
func (s *searchResults) Get(line int64) (searchResult, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if i, ok := s.lines.Get(line); ok {
		return s.results[i], true
	}

	return searchResult{}, false
}

// Next returns the search result for the nearest line after line,
// noninclusive, if one exists.
func (s *searchResults) Next(line int64) (searchResult, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, i, err := s.lines.NearestGreater(line)
	if err != nil {
		// Probably ErrNoSuchKey, aka none found.
		return searchResult{}, false
	}

	return s.results[i], true
}

// Prev returns the search result for the nearest line before line,
// noninclusive, if one exists.
func (s *searchResults) Prev(line int64) (searchResult, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Search for line - 1, since it may be equal.
	_, i, err := s.lines.NearestLessEqual(line - 1)
	if err != nil {
		// Probably ErrNoSuchKey, aka none found.
		return searchResult{}, false
	}

	return s.results[i], true
}

func (l *Lesser) search(s string) searchResults {
	reg, err := regexp.Compile(s)
	if err != nil {
		// TODO(prattmic): display a better error
		log.Printf("regexp failed to compile: %v", err)
		return NewSearchResults()
	}

	resultChan := make(chan searchResult, 100)

	searchLine := func(line int64) {
		r, err := l.src.SearchLine(reg, line)
		if err != nil {
			r = nil
		}

		resultChan <- searchResult{
			line:    line,
			matches: r,
			err:     err,
		}
	}

	nextLine := int64(1)
	// Spawn initial search goroutines
	for ; nextLine <= 5; nextLine++ {
		go searchLine(nextLine)
	}

	results := NewSearchResults()

	var count int64

	waitResult := func() searchResult {
		ret := <-resultChan
		count += 1

		// Only store results with matches.
		if len(ret.matches) > 0 {
			results.Add(ret)
		}

		return ret
	}

	// Collect results, start searching next lines until we start
	// hitting EOF.
	for {
		r := waitResult()

		// We started hitting errors on a previous line,
		// there is no reason to search later lines.
		if r.err != nil {
			break
		}

		go searchLine(nextLine)
		nextLine++
	}

	// Collect the remaing results.
	for count < nextLine-1 {
		waitResult()
	}

	return results
}

// statusBar renders the status bar.
// mu must be held on call.
func (l *Lesser) statusBar() {
	// The statusbar is just below the display.

	// Clear the statusbar
	for i := 0; i < l.size.x; i++ {
		termbox.SetCell(i, l.size.y, ' ', 0, 0)
	}

	switch l.mode {
	case ModeNormal:
		// Just a colon and a cursor
		termbox.SetCell(0, l.size.y, ':', 0, 0)
		termbox.SetCursor(1, l.size.y)
	case ModeSearchEntry:
		// / and search string
		termbox.SetCell(0, l.size.y, '/', 0, 0)
		for i, c := range l.regexp {
			termbox.SetCell(1+i, l.size.y, c, 0, 0)
		}
		termbox.SetCursor(1+len(l.regexp), l.size.y)
	}
}

// alignUp aligns n up to the next multiple of divisor.
func alignUp(n, divisor int) int {
	return n + (divisor - (n % divisor))
}

func (l *Lesser) refreshScreen() error {
	l.mu.Lock()
	defer l.mu.Unlock()

	for y := 0; y < l.size.y; y++ {
		buf := make([]byte, l.size.x)
		line := l.line + int64(y)

		_, err := l.src.ReadLine(buf, line)
		// EOF just means the line was shorter than the display.
		if err != nil && err != io.EOF {
			return err
		}

		highlight, ok := l.searchResults.Get(line)

		var displayColumn int
		for i, c := range buf {
			if displayColumn >= l.size.x {
				break
			}

			fg := termbox.ColorDefault
			bg := termbox.ColorDefault

			// Highlight matches
			if ok && highlight.matchesChar(i) {
				fg = termbox.ColorBlack
				bg = termbox.ColorWhite
			}

			if c == '\t' {
				// Tabs align the display up to the next
				// multiple of tabstop.
				next := alignUp(displayColumn, l.tabStop)

				// Clear the tab spaces
				for j := displayColumn; j < next; j++ {
					termbox.SetCell(j, y, ' ', 0, 0)
				}

				displayColumn = next
			} else {
				termbox.SetCell(displayColumn, y, rune(c), fg, bg)
				displayColumn += 1
			}
		}
	}

	l.statusBar()

	termbox.Flush()

	return nil
}

func (l *Lesser) Run() {
	go l.listenEvents()

	err := l.refreshScreen()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to refresh screen: %v\n", err)
		return
	}

	for {
		e := <-l.events

		switch e {
		case EventQuit:
			return
		case EventRefresh:
			err = l.refreshScreen()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Failed to refresh screen: %v\n", err)
				return
			}
		}
	}
}

func NewLesser(f *os.File, ts int) Lesser {
	x, y := termbox.Size()

	return Lesser{
		src:     lineio.NewLineReader(f),
		tabStop: ts,
		// Save one line for statusbar.
		size:   size{x: x, y: y - 1},
		line:   1,
		events: make(chan Event, 1),
		mode:   ModeNormal,
	}
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2014 The heketi 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 glusterfs

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"sync"
	"time"
)

type BrickNode struct {
	NodeId   string `json:"meta"`
	DeviceId string `json:"device"`
}
type BrickNodes []BrickNode

type RingOutput struct {
	Nodes     BrickNodes `json:"nodes"`
	Partition int        `json:"partition"`
}

type GlusterRing struct {
	db           *GlusterFSDB
	ringCreateCh chan bool
	lock         sync.RWMutex
}

func (g *GlusterRing) createServer() {

	for {
		// Wait for a command to start creating the ring
		select {
		case <-g.ringCreateCh:
			// This is not very efficient, but it works for now
			g.lock.Lock()
		}

		// Once we get the ring command.  Wait a maximum
		// of 5 seconds for more requests.  That way we only
		// do it once
		for created := false; !created; {
			timeout := time.After(time.Second * 5)
			select {
			case <-g.ringCreateCh:
				continue
			case <-timeout:
				err := g.createRing()

				// :TODO: Log
				if err != nil {
					fmt.Println(err)
				} else {
					fmt.Println("Success")
				}
				created = true
			}
		}

		g.lock.Unlock()
	}
}

func NewGlusterRing(db *GlusterFSDB) *GlusterRing {
	g := &GlusterRing{
		db:           db,
		ringCreateCh: make(chan bool),
	}

	go g.createServer()
	return g
}

func (g *GlusterRing) CreateRing() {
	g.ringCreateCh <- true
}

func (g *GlusterRing) createRing() error {

	os.Remove("heketi.builder")
	os.Remove("heketi.ring.gz")

	args := []string{
		"heketi.builder",
		"create",
		"16",
		"2",
		"1",
	}

	// Create new ring
	err := exec.Command("swift-ring-builder", args...).Run()
	if err != nil {
		return errors.New("Unable to create brick placement db")
	}

	err = g.db.Reader(func() error {
		// Add all devices
		for nodeid, node := range g.db.nodes {
			for devid, dev := range node.Info.Devices {
				args := []string{
					"heketi.builder",
					"add",
					fmt.Sprintf("r1z%v-%v:80/%v_%v",
						node.Info.Zone,
						node.Info.Name,
						devid,
						nodeid),
					fmt.Sprintf("%v", dev.Weight),
				}
				err := exec.Command("swift-ring-builder", args...).Run()
				if err != nil {
					return err
				}
			}
		}

		return nil
	})
	if err != nil {
		return nil
	}

	// Rebalance
	return exec.Command("swift-ring-builder", "heketi.builder", "rebalance").Run()
}

func (g *GlusterRing) GetNodes(brick_num int, id string) (BrickNodes, error) {

	var out bytes.Buffer
	var nodes RingOutput

	g.lock.RLock()
	defer g.lock.RUnlock()

	args := []string{
		fmt.Sprintf("%v", brick_num),
		id,
	}

	cmd := exec.Command("./ring.py", args...)
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return nil, err
	}

	if err := json.Unmarshal(out.Bytes(), &nodes); err != nil {
		return nil, err
	}

	return nodes.Nodes, nil
}
<commit_msg>Removed debug print<commit_after>//
// Copyright (c) 2014 The heketi 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 glusterfs

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"sync"
	"time"
)

type BrickNode struct {
	NodeId   string `json:"meta"`
	DeviceId string `json:"device"`
}
type BrickNodes []BrickNode

type RingOutput struct {
	Nodes     BrickNodes `json:"nodes"`
	Partition int        `json:"partition"`
}

type GlusterRing struct {
	db           *GlusterFSDB
	ringCreateCh chan bool
	lock         sync.RWMutex
}

func (g *GlusterRing) createServer() {

	for {
		// Wait for a command to start creating the ring
		select {
		case <-g.ringCreateCh:
			// This is not very efficient, but it works for now
			g.lock.Lock()
		}

		// Once we get the ring command.  Wait a maximum
		// of 5 seconds for more requests.  That way we only
		// do it once
		for created := false; !created; {
			timeout := time.After(time.Second * 5)
			select {
			case <-g.ringCreateCh:
				continue
			case <-timeout:
				err := g.createRing()

				// :TODO: Log
				if err != nil {
					// :TODO: May want to try again a few times if we fail.
					fmt.Println(err)
				}
				created = true
			}
		}

		g.lock.Unlock()
	}
}

func NewGlusterRing(db *GlusterFSDB) *GlusterRing {
	g := &GlusterRing{
		db:           db,
		ringCreateCh: make(chan bool),
	}

	go g.createServer()
	return g
}

func (g *GlusterRing) CreateRing() {
	g.ringCreateCh <- true
}

func (g *GlusterRing) createRing() error {

	os.Remove("heketi.builder")
	os.Remove("heketi.ring.gz")

	args := []string{
		"heketi.builder",
		"create",
		"16",
		"2",
		"1",
	}

	// Create new ring
	err := exec.Command("swift-ring-builder", args...).Run()
	if err != nil {
		return errors.New("Unable to create brick placement db")
	}

	err = g.db.Reader(func() error {
		// Add all devices
		for nodeid, node := range g.db.nodes {
			for devid, dev := range node.Info.Devices {
				args := []string{
					"heketi.builder",
					"add",
					fmt.Sprintf("r1z%v-%v:80/%v_%v",
						node.Info.Zone,
						node.Info.Name,
						devid,
						nodeid),
					fmt.Sprintf("%v", dev.Weight),
				}
				err := exec.Command("swift-ring-builder", args...).Run()
				if err != nil {
					return err
				}
			}
		}

		return nil
	})
	if err != nil {
		return nil
	}

	// Rebalance
	return exec.Command("swift-ring-builder", "heketi.builder", "rebalance").Run()
}

func (g *GlusterRing) GetNodes(brick_num int, id string) (BrickNodes, error) {

	var out bytes.Buffer
	var nodes RingOutput

	g.lock.RLock()
	defer g.lock.RUnlock()

	args := []string{
		fmt.Sprintf("%v", brick_num),
		id,
	}

	cmd := exec.Command("./ring.py", args...)
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return nil, err
	}

	if err := json.Unmarshal(out.Bytes(), &nodes); err != nil {
		return nil, err
	}

	return nodes.Nodes, nil
}
<|endoftext|>


<commit_before>// Copyright 2009 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 http_test

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	. "net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"strings"
	"testing"
)

func TestQuery(t *testing.T) {
	req := &Request{Method: "GET"}
	req.URL, _ = url.Parse("http://www.google.com/search?q=foo&q=bar")
	if q := req.FormValue("q"); q != "foo" {
		t.Errorf(`req.FormValue("q") = %q, want "foo"`, q)
	}
}

func TestPostQuery(t *testing.T) {
	req, _ := NewRequest("POST", "http://www.google.com/search?q=foo&q=bar&both=x",
		strings.NewReader("z=post&both=y"))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded; param=value")

	if q := req.FormValue("q"); q != "foo" {
		t.Errorf(`req.FormValue("q") = %q, want "foo"`, q)
	}
	if z := req.FormValue("z"); z != "post" {
		t.Errorf(`req.FormValue("z") = %q, want "post"`, z)
	}
	if both := req.Form["both"]; !reflect.DeepEqual(both, []string{"x", "y"}) {
		t.Errorf(`req.FormValue("both") = %q, want ["x", "y"]`, both)
	}
}

type stringMap map[string][]string
type parseContentTypeTest struct {
	contentType stringMap
}

var parseContentTypeTests = []parseContentTypeTest{
	{contentType: stringMap{"Content-Type": {"text/plain"}}},
	{contentType: stringMap{}}, // Non-existent keys are not placed. The value nil is illegal.
	{contentType: stringMap{"Content-Type": {"text/plain; boundary="}}},
	{
		contentType: stringMap{"Content-Type": {"application/unknown"}},
	},
}

func TestParseFormBadContentType(t *testing.T) {
	for i, test := range parseContentTypeTests {
		req := &Request{
			Method: "POST",
			Header: Header(test.contentType),
			Body:   ioutil.NopCloser(bytes.NewBufferString("body")),
		}
		err := req.ParseForm()
		if err == nil {
			t.Errorf("test %d should have returned error", i)
		}
	}
}

func TestMultipartReader(t *testing.T) {
	req := &Request{
		Method: "POST",
		Header: Header{"Content-Type": {`multipart/form-data; boundary="foo123"`}},
		Body:   ioutil.NopCloser(new(bytes.Buffer)),
	}
	multipart, err := req.MultipartReader()
	if multipart == nil {
		t.Errorf("expected multipart; error: %v", err)
	}

	req.Header = Header{"Content-Type": {"text/plain"}}
	multipart, err = req.MultipartReader()
	if multipart != nil {
		t.Errorf("unexpected multipart for text/plain")
	}
}

func TestRedirect(t *testing.T) {
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		switch r.URL.Path {
		case "/":
			w.Header().Set("Location", "/foo/")
			w.WriteHeader(StatusSeeOther)
		case "/foo/":
			fmt.Fprintf(w, "foo")
		default:
			w.WriteHeader(StatusBadRequest)
		}
	}))
	defer ts.Close()

	var end = regexp.MustCompile("/foo/$")
	r, err := Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	r.Body.Close()
	url := r.Request.URL.String()
	if r.StatusCode != 200 || !end.MatchString(url) {
		t.Fatalf("Get got status %d at %q, want 200 matching /foo/$", r.StatusCode, url)
	}
}

func TestSetBasicAuth(t *testing.T) {
	r, _ := NewRequest("GET", "http://example.com/", nil)
	r.SetBasicAuth("Aladdin", "open sesame")
	if g, e := r.Header.Get("Authorization"), "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="; g != e {
		t.Errorf("got header %q, want %q", g, e)
	}
}

func TestMultipartRequest(t *testing.T) {
	// Test that we can read the values and files of a
	// multipart request with FormValue and FormFile,
	// and that ParseMultipartForm can be called multiple times.
	req := newTestMultipartRequest(t)
	if err := req.ParseMultipartForm(25); err != nil {
		t.Fatal("ParseMultipartForm first call:", err)
	}
	defer req.MultipartForm.RemoveAll()
	validateTestMultipartContents(t, req, false)
	if err := req.ParseMultipartForm(25); err != nil {
		t.Fatal("ParseMultipartForm second call:", err)
	}
	validateTestMultipartContents(t, req, false)
}

func TestMultipartRequestAuto(t *testing.T) {
	// Test that FormValue and FormFile automatically invoke
	// ParseMultipartForm and return the right values.
	req := newTestMultipartRequest(t)
	defer func() {
		if req.MultipartForm != nil {
			req.MultipartForm.RemoveAll()
		}
	}()
	validateTestMultipartContents(t, req, true)
}

func TestEmptyMultipartRequest(t *testing.T) {
	// Test that FormValue and FormFile automatically invoke
	// ParseMultipartForm and return the right values.
	req, err := NewRequest("GET", "/", nil)
	if err != nil {
		t.Errorf("NewRequest err = %q", err)
	}
	testMissingFile(t, req)
}

func TestRequestMultipartCallOrder(t *testing.T) {
	req := newTestMultipartRequest(t)
	_, err := req.MultipartReader()
	if err != nil {
		t.Fatalf("MultipartReader: %v", err)
	}
	err = req.ParseMultipartForm(1024)
	if err == nil {
		t.Errorf("expected an error from ParseMultipartForm after call to MultipartReader")
	}
}

func testMissingFile(t *testing.T, req *Request) {
	f, fh, err := req.FormFile("missing")
	if f != nil {
		t.Errorf("FormFile file = %q, want nil", f)
	}
	if fh != nil {
		t.Errorf("FormFile file header = %q, want nil", fh)
	}
	if err != ErrMissingFile {
		t.Errorf("FormFile err = %q, want ErrMissingFile", err)
	}
}

func newTestMultipartRequest(t *testing.T) *Request {
	b := bytes.NewBufferString(strings.Replace(message, "\n", "\r\n", -1))
	req, err := NewRequest("POST", "/", b)
	if err != nil {
		t.Fatal("NewRequest:", err)
	}
	ctype := fmt.Sprintf(`multipart/form-data; boundary="%s"`, boundary)
	req.Header.Set("Content-type", ctype)
	return req
}

func validateTestMultipartContents(t *testing.T, req *Request, allMem bool) {
	if g, e := req.FormValue("texta"), textaValue; g != e {
		t.Errorf("texta value = %q, want %q", g, e)
	}
	if g, e := req.FormValue("texta"), textaValue; g != e {
		t.Errorf("texta value = %q, want %q", g, e)
	}
	if g := req.FormValue("missing"); g != "" {
		t.Errorf("missing value = %q, want empty string", g)
	}

	assertMem := func(n string, fd multipart.File) {
		if _, ok := fd.(*os.File); ok {
			t.Error(n, " is *os.File, should not be")
		}
	}
	fd := testMultipartFile(t, req, "filea", "filea.txt", fileaContents)
	assertMem("filea", fd)
	fd = testMultipartFile(t, req, "fileb", "fileb.txt", filebContents)
	if allMem {
		assertMem("fileb", fd)
	} else {
		if _, ok := fd.(*os.File); !ok {
			t.Errorf("fileb has unexpected underlying type %T", fd)
		}
	}

	testMissingFile(t, req)
}

func testMultipartFile(t *testing.T, req *Request, key, expectFilename, expectContent string) multipart.File {
	f, fh, err := req.FormFile(key)
	if err != nil {
		t.Fatalf("FormFile(%q): %q", key, err)
	}
	if fh.Filename != expectFilename {
		t.Errorf("filename = %q, want %q", fh.Filename, expectFilename)
	}
	var b bytes.Buffer
	_, err = io.Copy(&b, f)
	if err != nil {
		t.Fatal("copying contents:", err)
	}
	if g := b.String(); g != expectContent {
		t.Errorf("contents = %q, want %q", g, expectContent)
	}
	return f
}

const (
	fileaContents = "This is a test file."
	filebContents = "Another test file."
	textaValue    = "foo"
	textbValue    = "bar"
	boundary      = `MyBoundary`
)

const message = `
--MyBoundary
Content-Disposition: form-data; name="filea"; filename="filea.txt"
Content-Type: text/plain

` + fileaContents + `
--MyBoundary
Content-Disposition: form-data; name="fileb"; filename="fileb.txt"
Content-Type: text/plain

` + filebContents + `
--MyBoundary
Content-Disposition: form-data; name="texta"

` + textaValue + `
--MyBoundary
Content-Disposition: form-data; name="textb"

` + textbValue + `
--MyBoundary--
`
<commit_msg>net/http: test both texta and textb values, not texta twice<commit_after>// Copyright 2009 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 http_test

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	. "net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"strings"
	"testing"
)

func TestQuery(t *testing.T) {
	req := &Request{Method: "GET"}
	req.URL, _ = url.Parse("http://www.google.com/search?q=foo&q=bar")
	if q := req.FormValue("q"); q != "foo" {
		t.Errorf(`req.FormValue("q") = %q, want "foo"`, q)
	}
}

func TestPostQuery(t *testing.T) {
	req, _ := NewRequest("POST", "http://www.google.com/search?q=foo&q=bar&both=x",
		strings.NewReader("z=post&both=y"))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded; param=value")

	if q := req.FormValue("q"); q != "foo" {
		t.Errorf(`req.FormValue("q") = %q, want "foo"`, q)
	}
	if z := req.FormValue("z"); z != "post" {
		t.Errorf(`req.FormValue("z") = %q, want "post"`, z)
	}
	if both := req.Form["both"]; !reflect.DeepEqual(both, []string{"x", "y"}) {
		t.Errorf(`req.FormValue("both") = %q, want ["x", "y"]`, both)
	}
}

type stringMap map[string][]string
type parseContentTypeTest struct {
	contentType stringMap
}

var parseContentTypeTests = []parseContentTypeTest{
	{contentType: stringMap{"Content-Type": {"text/plain"}}},
	{contentType: stringMap{}}, // Non-existent keys are not placed. The value nil is illegal.
	{contentType: stringMap{"Content-Type": {"text/plain; boundary="}}},
	{
		contentType: stringMap{"Content-Type": {"application/unknown"}},
	},
}

func TestParseFormBadContentType(t *testing.T) {
	for i, test := range parseContentTypeTests {
		req := &Request{
			Method: "POST",
			Header: Header(test.contentType),
			Body:   ioutil.NopCloser(bytes.NewBufferString("body")),
		}
		err := req.ParseForm()
		if err == nil {
			t.Errorf("test %d should have returned error", i)
		}
	}
}

func TestMultipartReader(t *testing.T) {
	req := &Request{
		Method: "POST",
		Header: Header{"Content-Type": {`multipart/form-data; boundary="foo123"`}},
		Body:   ioutil.NopCloser(new(bytes.Buffer)),
	}
	multipart, err := req.MultipartReader()
	if multipart == nil {
		t.Errorf("expected multipart; error: %v", err)
	}

	req.Header = Header{"Content-Type": {"text/plain"}}
	multipart, err = req.MultipartReader()
	if multipart != nil {
		t.Errorf("unexpected multipart for text/plain")
	}
}

func TestRedirect(t *testing.T) {
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		switch r.URL.Path {
		case "/":
			w.Header().Set("Location", "/foo/")
			w.WriteHeader(StatusSeeOther)
		case "/foo/":
			fmt.Fprintf(w, "foo")
		default:
			w.WriteHeader(StatusBadRequest)
		}
	}))
	defer ts.Close()

	var end = regexp.MustCompile("/foo/$")
	r, err := Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	r.Body.Close()
	url := r.Request.URL.String()
	if r.StatusCode != 200 || !end.MatchString(url) {
		t.Fatalf("Get got status %d at %q, want 200 matching /foo/$", r.StatusCode, url)
	}
}

func TestSetBasicAuth(t *testing.T) {
	r, _ := NewRequest("GET", "http://example.com/", nil)
	r.SetBasicAuth("Aladdin", "open sesame")
	if g, e := r.Header.Get("Authorization"), "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="; g != e {
		t.Errorf("got header %q, want %q", g, e)
	}
}

func TestMultipartRequest(t *testing.T) {
	// Test that we can read the values and files of a
	// multipart request with FormValue and FormFile,
	// and that ParseMultipartForm can be called multiple times.
	req := newTestMultipartRequest(t)
	if err := req.ParseMultipartForm(25); err != nil {
		t.Fatal("ParseMultipartForm first call:", err)
	}
	defer req.MultipartForm.RemoveAll()
	validateTestMultipartContents(t, req, false)
	if err := req.ParseMultipartForm(25); err != nil {
		t.Fatal("ParseMultipartForm second call:", err)
	}
	validateTestMultipartContents(t, req, false)
}

func TestMultipartRequestAuto(t *testing.T) {
	// Test that FormValue and FormFile automatically invoke
	// ParseMultipartForm and return the right values.
	req := newTestMultipartRequest(t)
	defer func() {
		if req.MultipartForm != nil {
			req.MultipartForm.RemoveAll()
		}
	}()
	validateTestMultipartContents(t, req, true)
}

func TestEmptyMultipartRequest(t *testing.T) {
	// Test that FormValue and FormFile automatically invoke
	// ParseMultipartForm and return the right values.
	req, err := NewRequest("GET", "/", nil)
	if err != nil {
		t.Errorf("NewRequest err = %q", err)
	}
	testMissingFile(t, req)
}

func TestRequestMultipartCallOrder(t *testing.T) {
	req := newTestMultipartRequest(t)
	_, err := req.MultipartReader()
	if err != nil {
		t.Fatalf("MultipartReader: %v", err)
	}
	err = req.ParseMultipartForm(1024)
	if err == nil {
		t.Errorf("expected an error from ParseMultipartForm after call to MultipartReader")
	}
}

func testMissingFile(t *testing.T, req *Request) {
	f, fh, err := req.FormFile("missing")
	if f != nil {
		t.Errorf("FormFile file = %q, want nil", f)
	}
	if fh != nil {
		t.Errorf("FormFile file header = %q, want nil", fh)
	}
	if err != ErrMissingFile {
		t.Errorf("FormFile err = %q, want ErrMissingFile", err)
	}
}

func newTestMultipartRequest(t *testing.T) *Request {
	b := bytes.NewBufferString(strings.Replace(message, "\n", "\r\n", -1))
	req, err := NewRequest("POST", "/", b)
	if err != nil {
		t.Fatal("NewRequest:", err)
	}
	ctype := fmt.Sprintf(`multipart/form-data; boundary="%s"`, boundary)
	req.Header.Set("Content-type", ctype)
	return req
}

func validateTestMultipartContents(t *testing.T, req *Request, allMem bool) {
	if g, e := req.FormValue("texta"), textaValue; g != e {
		t.Errorf("texta value = %q, want %q", g, e)
	}
	if g, e := req.FormValue("textb"), textbValue; g != e {
		t.Errorf("textb value = %q, want %q", g, e)
	}
	if g := req.FormValue("missing"); g != "" {
		t.Errorf("missing value = %q, want empty string", g)
	}

	assertMem := func(n string, fd multipart.File) {
		if _, ok := fd.(*os.File); ok {
			t.Error(n, " is *os.File, should not be")
		}
	}
	fd := testMultipartFile(t, req, "filea", "filea.txt", fileaContents)
	assertMem("filea", fd)
	fd = testMultipartFile(t, req, "fileb", "fileb.txt", filebContents)
	if allMem {
		assertMem("fileb", fd)
	} else {
		if _, ok := fd.(*os.File); !ok {
			t.Errorf("fileb has unexpected underlying type %T", fd)
		}
	}

	testMissingFile(t, req)
}

func testMultipartFile(t *testing.T, req *Request, key, expectFilename, expectContent string) multipart.File {
	f, fh, err := req.FormFile(key)
	if err != nil {
		t.Fatalf("FormFile(%q): %q", key, err)
	}
	if fh.Filename != expectFilename {
		t.Errorf("filename = %q, want %q", fh.Filename, expectFilename)
	}
	var b bytes.Buffer
	_, err = io.Copy(&b, f)
	if err != nil {
		t.Fatal("copying contents:", err)
	}
	if g := b.String(); g != expectContent {
		t.Errorf("contents = %q, want %q", g, expectContent)
	}
	return f
}

const (
	fileaContents = "This is a test file."
	filebContents = "Another test file."
	textaValue    = "foo"
	textbValue    = "bar"
	boundary      = `MyBoundary`
)

const message = `
--MyBoundary
Content-Disposition: form-data; name="filea"; filename="filea.txt"
Content-Type: text/plain

` + fileaContents + `
--MyBoundary
Content-Disposition: form-data; name="fileb"; filename="fileb.txt"
Content-Type: text/plain

` + filebContents + `
--MyBoundary
Content-Disposition: form-data; name="texta"

` + textaValue + `
--MyBoundary
Content-Disposition: form-data; name="textb"

` + textbValue + `
--MyBoundary--
`
<|endoftext|>