<commit_before>package fakes

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/chunker"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/rexec"
	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	rc "github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
	regrpc "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	repb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	log "github.com/golang/glog"
	tspb "github.com/golang/protobuf/ptypes/timestamp"
	bsgrpc "google.golang.org/genproto/googleapis/bytestream"
)

// Server is a configurable fake in-process RBE server for use in integration tests.
type Server struct {
	Exec        *Exec
	CAS         *CAS
	ActionCache *ActionCache
	listener    net.Listener
	srv         *grpc.Server
}

// NewServer creates a server that is ready to accept requests.
func NewServer(t testing.TB) (s *Server, err error) {
	cas := NewCAS()
	ac := NewActionCache()
	s = &Server{Exec: NewExec(t, ac, cas), CAS: cas, ActionCache: ac}
	s.listener, err = net.Listen("tcp", ":0")
	if err != nil {
		return nil, err
	}
	s.srv = grpc.NewServer()
	bsgrpc.RegisterByteStreamServer(s.srv, s.CAS)
	regrpc.RegisterContentAddressableStorageServer(s.srv, s.CAS)
	regrpc.RegisterActionCacheServer(s.srv, s.ActionCache)
	regrpc.RegisterCapabilitiesServer(s.srv, s.Exec)
	regrpc.RegisterExecutionServer(s.srv, s.Exec)
	go s.srv.Serve(s.listener)
	return s, nil
}

// Clear clears the fake results.
func (s *Server) Clear() {
	s.CAS.Clear()
	s.ActionCache.Clear()
	s.Exec.Clear()
}

// Stop shuts down the in process server.
func (s *Server) Stop() {
	s.listener.Close()
	s.srv.Stop()
}

// NewTestClient returns a new in-process Client connected to this server.
func (s *Server) NewTestClient(ctx context.Context) (*rc.Client, error) {
	return rc.NewClient(ctx, "instance", s.dialParams())
}

// NewClientConn returns a gRPC client connction to the server.
func (s *Server) NewClientConn(ctx context.Context) (*grpc.ClientConn, error) {
	p := s.dialParams()
	return client.Dial(ctx, p.Service, p)
}

func (s *Server) dialParams() rc.DialParams {
	return rc.DialParams{
		Service:    s.listener.Addr().String(),
		NoSecurity: true,
	}
}

// TestEnv is a wrapper for convenient integration tests of remote execution.
type TestEnv struct {
	Client   *rexec.Client
	Server   *Server
	ExecRoot string
	t        testing.TB
}

// NewTestEnv initializes a TestEnv containing a fake server, a client connected to it,
// and a temporary directory used as execution root for inputs and outputs.
// It returns the new env and a cleanup function that should be called in the end of the test.
func NewTestEnv(t testing.TB) (*TestEnv, func()) {
	t.Helper()
	// Set up temp directory.
	execRoot, err := ioutil.TempDir("", strings.ReplaceAll(t.Name(), string(filepath.Separator), "_"))
	if err != nil {
		t.Fatalf("failed to make temp dir: %v", err)
	}
	// Set up the fake.
	s, err := NewServer(t)
	if err != nil {
		t.Fatalf("Error starting fake server: %v", err)
	}
	grpcClient, err := s.NewTestClient(context.Background())
	if err != nil {
		t.Fatalf("Error connecting to server: %v", err)
	}
	return &TestEnv{
			Client: &rexec.Client{
				FileMetadataCache: filemetadata.NewNoopCache(),
				GrpcClient:        grpcClient,
			},
			Server:   s,
			ExecRoot: execRoot,
			t:        t,
		}, func() {
			grpcClient.Close()
			s.Stop()
			os.RemoveAll(execRoot)
		}
}

func timeToProto(t time.Time) *tspb.Timestamp {
	if t.IsZero() {
		return nil
	}
	ts, err := ptypes.TimestampProto(t)
	if err != nil {
		log.Warningf("Unable to convert time to Timestamp: %v", err.Error())
		return nil
	}
	return ts
}

// Set sets up the fake to return the given result on the given command execution.
// It is not possible to make the fake result in a LocalErrorResultStatus or an InterruptedResultStatus.
func (e *TestEnv) Set(cmd *command.Command, opt *command.ExecutionOptions, res *command.Result, opts ...option) (cmdDg, acDg digest.Digest) {
	e.t.Helper()
	cmd.FillDefaultFieldValues()
	if err := cmd.Validate(); err != nil {
		e.t.Fatalf("command validation failed: %v", err)
	}

	t, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:05Z")
	ar := &repb.ActionResult{
		ExitCode: int32(res.ExitCode),
		ExecutionMetadata: &repb.ExecutedActionMetadata{
			QueuedTimestamp:                timeToProto(t.Add(time.Millisecond)),
			WorkerStartTimestamp:           timeToProto(t.Add(2 * time.Millisecond)),
			WorkerCompletedTimestamp:       timeToProto(t.Add(3 * time.Millisecond)),
			InputFetchStartTimestamp:       timeToProto(t.Add(4 * time.Millisecond)),
			InputFetchCompletedTimestamp:   timeToProto(t.Add(5 * time.Millisecond)),
			ExecutionStartTimestamp:        timeToProto(t.Add(6 * time.Millisecond)),
			ExecutionCompletedTimestamp:    timeToProto(t.Add(7 * time.Millisecond)),
			OutputUploadStartTimestamp:     timeToProto(t.Add(8 * time.Millisecond)),
			OutputUploadCompletedTimestamp: timeToProto(t.Add(9 * time.Millisecond)),
		},
	}
	for _, o := range opts {
		if err := o.Apply(ar, e.Server, e.ExecRoot); err != nil {
			e.t.Fatalf("error applying option %+v: %v", o, err)
		}
	}

	execRoot, workingDir, remoteWorkingDir := cmd.ExecRoot, cmd.WorkingDir, cmd.RemoteWorkingDir
	root, inputs, _, err := e.Client.GrpcClient.ComputeMerkleTree(execRoot, workingDir, remoteWorkingDir, cmd.InputSpec, e.Client.FileMetadataCache)
	if err != nil {
		e.t.Fatalf("error building input tree in fake setup: %v", err)
		return digest.Empty, digest.Empty
	}
	for _, inp := range inputs {
		ch, err := chunker.New(inp, false, int(e.Client.GrpcClient.ChunkMaxSize))
		if err != nil {
			e.t.Fatalf("error getting data from input entry: %v", err)
		}
		bytes, err := ch.FullData()
		if err != nil {
			e.t.Fatalf("error getting data from input chunker: %v", err)
		}
		e.Server.CAS.Put(bytes)
	}

	cmdPb := cmd.ToREProto(false)
	bytes, err := proto.Marshal(cmdPb)
	if err != nil {
		e.t.Fatalf("error inserting command digest blob into CAS %v", err)
	}
	e.Server.CAS.Put(bytes)

	cmdDg = digest.TestNewFromMessage(cmdPb)
	ac := &repb.Action{
		CommandDigest:   cmdDg.ToProto(),
		InputRootDigest: root.ToProto(),
		DoNotCache:      opt.DoNotCache,
	}
	if cmd.Timeout > 0 {
		ac.Timeout = ptypes.DurationProto(cmd.Timeout)
	}
	acDg = digest.TestNewFromMessage(ac)

	bytes, err = proto.Marshal(ac)
	if err != nil {
		e.t.Fatalf("error inserting action digest blob into CAS %v", err)
	}
	e.Server.CAS.Put(bytes)

	e.Server.Exec.adg = acDg
	e.Server.Exec.ActionResult = ar
	switch res.Status {
	case command.TimeoutResultStatus:
		if res.Err == nil {
			e.Server.Exec.Status = status.New(codes.DeadlineExceeded, "timeout")
		} else {
			e.Server.Exec.Status = status.New(codes.DeadlineExceeded, res.Err.Error())
		}
	case command.RemoteErrorResultStatus:
		st, ok := status.FromError(res.Err)
		if !ok {
			if res.Err == nil {
				st = status.New(codes.Internal, "remote error")
			} else {
				st = status.New(codes.Internal, res.Err.Error())
			}
		}
		e.Server.Exec.Status = st
	case command.CacheHitResultStatus:
		if !e.Server.Exec.Cached { // Assume the user means in this case the actual ActionCache should miss.
			e.Server.ActionCache.Put(acDg, ar)
		}
	}
	return cmdDg, acDg
}

type option interface {
	Apply(*repb.ActionResult, *Server, string) error
}

// InputFile to be made available to the fake action.
type InputFile struct {
	Path     string
	Contents string
}

// Apply creates a file in the execroot with the given content
// and also inserts the file blob into CAS.
func (f *InputFile) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	bytes := []byte(f.Contents)
	if err := os.MkdirAll(filepath.Join(execRoot, filepath.Dir(f.Path)), os.ModePerm); err != nil {
		return fmt.Errorf("failed to create input dir %v: %v", filepath.Dir(f.Path), err)
	}
	err := ioutil.WriteFile(filepath.Join(execRoot, f.Path), nil, 0755)
	if err != nil {
		return fmt.Errorf("failed to setup file %v under temp exec root %v: %v", f.Path, execRoot, err)
	}
	s.CAS.Put(bytes)
	return nil
}

// OutputFile is to be added as an output of the fake action.
type OutputFile struct {
	Path     string
	Contents string
}

// Apply puts the file in the fake CAS and the given ActionResult.
func (f *OutputFile) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	bytes := []byte(f.Contents)
	s.Exec.OutputBlobs = append(s.Exec.OutputBlobs, bytes)
	dg := s.CAS.Put(bytes)
	ac.OutputFiles = append(ac.OutputFiles, &repb.OutputFile{Path: f.Path, Digest: dg.ToProto()})
	return nil
}

// StdOut is to be added as an output of the fake action.
type StdOut string

// Apply puts the action stdout in the fake CAS and the given ActionResult.
func (o StdOut) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	bytes := []byte(o)
	s.Exec.OutputBlobs = append(s.Exec.OutputBlobs, bytes)
	dg := s.CAS.Put(bytes)
	ac.StdoutDigest = dg.ToProto()
	return nil
}

// StdOutRaw is to be added as a raw output of the fake action.
type StdOutRaw string

// Apply puts the action stdout as raw bytes in the given ActionResult.
func (o StdOutRaw) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	ac.StdoutRaw = []byte(o)
	return nil
}

// StdErr is to be added as an output of the fake action.
type StdErr string

// Apply puts the action stderr in the fake CAS and the given ActionResult.
func (o StdErr) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	bytes := []byte(o)
	s.Exec.OutputBlobs = append(s.Exec.OutputBlobs, bytes)
	dg := s.CAS.Put(bytes)
	ac.StderrDigest = dg.ToProto()
	return nil
}

// StdErrRaw is to be added as a raw output of the fake action.
type StdErrRaw string

// Apply puts the action stderr as raw bytes in the given ActionResult.
func (o StdErrRaw) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	ac.StderrRaw = []byte(o)
	return nil
}

// ExecutionCacheHit of true will cause the ActionResult to be returned as a cache hit during
// fake execution.
type ExecutionCacheHit bool

// Apply on true will cause the ActionResult to be returned as a cache hit during fake execution.
func (c ExecutionCacheHit) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	s.Exec.Cached = bool(c)
	return nil
}
<commit_msg>Add support for output directories in fake server. (#397)<commit_after>package fakes

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/chunker"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/rexec"
	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	rc "github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
	regrpc "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	repb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	log "github.com/golang/glog"
	tspb "github.com/golang/protobuf/ptypes/timestamp"
	bsgrpc "google.golang.org/genproto/googleapis/bytestream"
)

// Server is a configurable fake in-process RBE server for use in integration tests.
type Server struct {
	Exec        *Exec
	CAS         *CAS
	ActionCache *ActionCache
	listener    net.Listener
	srv         *grpc.Server
}

// NewServer creates a server that is ready to accept requests.
func NewServer(t testing.TB) (s *Server, err error) {
	cas := NewCAS()
	ac := NewActionCache()
	s = &Server{Exec: NewExec(t, ac, cas), CAS: cas, ActionCache: ac}
	s.listener, err = net.Listen("tcp", ":0")
	if err != nil {
		return nil, err
	}
	s.srv = grpc.NewServer()
	bsgrpc.RegisterByteStreamServer(s.srv, s.CAS)
	regrpc.RegisterContentAddressableStorageServer(s.srv, s.CAS)
	regrpc.RegisterActionCacheServer(s.srv, s.ActionCache)
	regrpc.RegisterCapabilitiesServer(s.srv, s.Exec)
	regrpc.RegisterExecutionServer(s.srv, s.Exec)
	go s.srv.Serve(s.listener)
	return s, nil
}

// Clear clears the fake results.
func (s *Server) Clear() {
	s.CAS.Clear()
	s.ActionCache.Clear()
	s.Exec.Clear()
}

// Stop shuts down the in process server.
func (s *Server) Stop() {
	s.listener.Close()
	s.srv.Stop()
}

// NewTestClient returns a new in-process Client connected to this server.
func (s *Server) NewTestClient(ctx context.Context) (*rc.Client, error) {
	return rc.NewClient(ctx, "instance", s.dialParams())
}

// NewClientConn returns a gRPC client connction to the server.
func (s *Server) NewClientConn(ctx context.Context) (*grpc.ClientConn, error) {
	p := s.dialParams()
	return client.Dial(ctx, p.Service, p)
}

func (s *Server) dialParams() rc.DialParams {
	return rc.DialParams{
		Service:    s.listener.Addr().String(),
		NoSecurity: true,
	}
}

// TestEnv is a wrapper for convenient integration tests of remote execution.
type TestEnv struct {
	Client   *rexec.Client
	Server   *Server
	ExecRoot string
	t        testing.TB
}

// NewTestEnv initializes a TestEnv containing a fake server, a client connected to it,
// and a temporary directory used as execution root for inputs and outputs.
// It returns the new env and a cleanup function that should be called in the end of the test.
func NewTestEnv(t testing.TB) (*TestEnv, func()) {
	t.Helper()
	// Set up temp directory.
	execRoot, err := ioutil.TempDir("", strings.ReplaceAll(t.Name(), string(filepath.Separator), "_"))
	if err != nil {
		t.Fatalf("failed to make temp dir: %v", err)
	}
	// Set up the fake.
	s, err := NewServer(t)
	if err != nil {
		t.Fatalf("Error starting fake server: %v", err)
	}
	grpcClient, err := s.NewTestClient(context.Background())
	if err != nil {
		t.Fatalf("Error connecting to server: %v", err)
	}
	return &TestEnv{
			Client: &rexec.Client{
				FileMetadataCache: filemetadata.NewNoopCache(),
				GrpcClient:        grpcClient,
			},
			Server:   s,
			ExecRoot: execRoot,
			t:        t,
		}, func() {
			grpcClient.Close()
			s.Stop()
			os.RemoveAll(execRoot)
		}
}

func timeToProto(t time.Time) *tspb.Timestamp {
	if t.IsZero() {
		return nil
	}
	ts, err := ptypes.TimestampProto(t)
	if err != nil {
		log.Warningf("Unable to convert time to Timestamp: %v", err.Error())
		return nil
	}
	return ts
}

// Set sets up the fake to return the given result on the given command execution.
// It is not possible to make the fake result in a LocalErrorResultStatus or an InterruptedResultStatus.
func (e *TestEnv) Set(cmd *command.Command, opt *command.ExecutionOptions, res *command.Result, opts ...option) (cmdDg, acDg digest.Digest) {
	e.t.Helper()
	cmd.FillDefaultFieldValues()
	if err := cmd.Validate(); err != nil {
		e.t.Fatalf("command validation failed: %v", err)
	}

	t, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:05Z")
	ar := &repb.ActionResult{
		ExitCode: int32(res.ExitCode),
		ExecutionMetadata: &repb.ExecutedActionMetadata{
			QueuedTimestamp:                timeToProto(t.Add(time.Millisecond)),
			WorkerStartTimestamp:           timeToProto(t.Add(2 * time.Millisecond)),
			WorkerCompletedTimestamp:       timeToProto(t.Add(3 * time.Millisecond)),
			InputFetchStartTimestamp:       timeToProto(t.Add(4 * time.Millisecond)),
			InputFetchCompletedTimestamp:   timeToProto(t.Add(5 * time.Millisecond)),
			ExecutionStartTimestamp:        timeToProto(t.Add(6 * time.Millisecond)),
			ExecutionCompletedTimestamp:    timeToProto(t.Add(7 * time.Millisecond)),
			OutputUploadStartTimestamp:     timeToProto(t.Add(8 * time.Millisecond)),
			OutputUploadCompletedTimestamp: timeToProto(t.Add(9 * time.Millisecond)),
		},
	}
	for _, o := range opts {
		if err := o.Apply(ar, e.Server, e.ExecRoot); err != nil {
			e.t.Fatalf("error applying option %+v: %v", o, err)
		}
	}

	execRoot, workingDir, remoteWorkingDir := cmd.ExecRoot, cmd.WorkingDir, cmd.RemoteWorkingDir
	root, inputs, _, err := e.Client.GrpcClient.ComputeMerkleTree(execRoot, workingDir, remoteWorkingDir, cmd.InputSpec, e.Client.FileMetadataCache)
	if err != nil {
		e.t.Fatalf("error building input tree in fake setup: %v", err)
		return digest.Empty, digest.Empty
	}
	for _, inp := range inputs {
		ch, err := chunker.New(inp, false, int(e.Client.GrpcClient.ChunkMaxSize))
		if err != nil {
			e.t.Fatalf("error getting data from input entry: %v", err)
		}
		bytes, err := ch.FullData()
		if err != nil {
			e.t.Fatalf("error getting data from input chunker: %v", err)
		}
		e.Server.CAS.Put(bytes)
	}

	cmdPb := cmd.ToREProto(false)
	bytes, err := proto.Marshal(cmdPb)
	if err != nil {
		e.t.Fatalf("error inserting command digest blob into CAS %v", err)
	}
	e.Server.CAS.Put(bytes)

	cmdDg = digest.TestNewFromMessage(cmdPb)
	ac := &repb.Action{
		CommandDigest:   cmdDg.ToProto(),
		InputRootDigest: root.ToProto(),
		DoNotCache:      opt.DoNotCache,
	}
	if cmd.Timeout > 0 {
		ac.Timeout = ptypes.DurationProto(cmd.Timeout)
	}
	acDg = digest.TestNewFromMessage(ac)

	bytes, err = proto.Marshal(ac)
	if err != nil {
		e.t.Fatalf("error inserting action digest blob into CAS %v", err)
	}
	e.Server.CAS.Put(bytes)

	e.Server.Exec.adg = acDg
	e.Server.Exec.ActionResult = ar
	switch res.Status {
	case command.TimeoutResultStatus:
		if res.Err == nil {
			e.Server.Exec.Status = status.New(codes.DeadlineExceeded, "timeout")
		} else {
			e.Server.Exec.Status = status.New(codes.DeadlineExceeded, res.Err.Error())
		}
	case command.RemoteErrorResultStatus:
		st, ok := status.FromError(res.Err)
		if !ok {
			if res.Err == nil {
				st = status.New(codes.Internal, "remote error")
			} else {
				st = status.New(codes.Internal, res.Err.Error())
			}
		}
		e.Server.Exec.Status = st
	case command.CacheHitResultStatus:
		if !e.Server.Exec.Cached { // Assume the user means in this case the actual ActionCache should miss.
			e.Server.ActionCache.Put(acDg, ar)
		}
	}
	return cmdDg, acDg
}

type option interface {
	Apply(*repb.ActionResult, *Server, string) error
}

// InputFile to be made available to the fake action.
type InputFile struct {
	Path     string
	Contents string
}

// Apply creates a file in the execroot with the given content
// and also inserts the file blob into CAS.
func (f *InputFile) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	bytes := []byte(f.Contents)
	if err := os.MkdirAll(filepath.Join(execRoot, filepath.Dir(f.Path)), os.ModePerm); err != nil {
		return fmt.Errorf("failed to create input dir %v: %v", filepath.Dir(f.Path), err)
	}
	err := ioutil.WriteFile(filepath.Join(execRoot, f.Path), nil, 0755)
	if err != nil {
		return fmt.Errorf("failed to setup file %v under temp exec root %v: %v", f.Path, execRoot, err)
	}
	s.CAS.Put(bytes)
	return nil
}

// OutputFile is to be added as an output of the fake action.
type OutputFile struct {
	Path     string
	Contents string
}

// Apply puts the file in the fake CAS and the given ActionResult.
func (f *OutputFile) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	bytes := []byte(f.Contents)
	s.Exec.OutputBlobs = append(s.Exec.OutputBlobs, bytes)
	dg := s.CAS.Put(bytes)
	ac.OutputFiles = append(ac.OutputFiles, &repb.OutputFile{Path: f.Path, Digest: dg.ToProto()})
	return nil
}

// OutputDir is to be added as an output of the fake action.
type OutputDir struct {
	Path string
}

// Apply puts the file in the fake CAS and the given ActionResult.
func (d *OutputDir) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	root, ch, err := BuildDir(d.Path, s, execRoot)
	if err != nil {
		return fmt.Errorf("failed to build directory tree: %v", err)
	}
	tr := &repb.Tree{
		Root:     root,
		Children: ch,
	}
	treeBlob, err := proto.Marshal(tr)
	if err != nil {
		return fmt.Errorf("failed to marshal tree: %v", err)
	}
	treeDigest := s.CAS.Put(treeBlob)
	ac.OutputDirectories = append(ac.OutputDirectories, &repb.OutputDirectory{Path: d.Path, TreeDigest: treeDigest.ToProto()})
	return nil
}

// Builds the directory tree by recursively iterating through the directory.
// This is similar to tree.go ComputeMerkleTree.
func BuildDir(path string, s *Server, execRoot string) (root *repb.Directory, childDir []*repb.Directory, err error) {
	res := &repb.Directory{}
	ch := []*repb.Directory{}

	files, err := ioutil.ReadDir(filepath.Join(execRoot, path))
	if err != nil {
		return nil, nil, fmt.Errorf("failed read directory: %v", err)
	}

	for _, file := range files {
		fn := file.Name()
		fp := filepath.Join(execRoot, path, fn)
		if file.IsDir() {
			root, _, err := BuildDir(fp, s, execRoot)
			if err != nil {
				return nil, nil, fmt.Errorf("failed to build directory tree: %v", err)
			}
			res.Directories = append(res.Directories, &repb.DirectoryNode{Name: fn, Digest: digest.TestNewFromMessage(root).ToProto()})
			ch = append(ch, root)
		} else {
			content, err := ioutil.ReadFile(fp)
			if err != nil {
				return nil, nil, fmt.Errorf("failed to read file: %v", err)
			}
			dg := s.CAS.Put(content)
			res.Files = append(res.Files, &repb.FileNode{Name: fn, Digest: dg.ToProto()})
		}
	}
	return res, ch, nil
}

// StdOut is to be added as an output of the fake action.
type StdOut string

// Apply puts the action stdout in the fake CAS and the given ActionResult.
func (o StdOut) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	bytes := []byte(o)
	s.Exec.OutputBlobs = append(s.Exec.OutputBlobs, bytes)
	dg := s.CAS.Put(bytes)
	ac.StdoutDigest = dg.ToProto()
	return nil
}

// StdOutRaw is to be added as a raw output of the fake action.
type StdOutRaw string

// Apply puts the action stdout as raw bytes in the given ActionResult.
func (o StdOutRaw) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	ac.StdoutRaw = []byte(o)
	return nil
}

// StdErr is to be added as an output of the fake action.
type StdErr string

// Apply puts the action stderr in the fake CAS and the given ActionResult.
func (o StdErr) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	bytes := []byte(o)
	s.Exec.OutputBlobs = append(s.Exec.OutputBlobs, bytes)
	dg := s.CAS.Put(bytes)
	ac.StderrDigest = dg.ToProto()
	return nil
}

// StdErrRaw is to be added as a raw output of the fake action.
type StdErrRaw string

// Apply puts the action stderr as raw bytes in the given ActionResult.
func (o StdErrRaw) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	ac.StderrRaw = []byte(o)
	return nil
}

// ExecutionCacheHit of true will cause the ActionResult to be returned as a cache hit during
// fake execution.
type ExecutionCacheHit bool

// Apply on true will cause the ActionResult to be returned as a cache hit during fake execution.
func (c ExecutionCacheHit) Apply(ac *repb.ActionResult, s *Server, execRoot string) error {
	s.Exec.Cached = bool(c)
	return nil
}
<|endoftext|>


<commit_before>package registry

import (
	"encoding/json"
	"fmt"
	"os"
	"path"
	"strings"

	"code.google.com/p/go-uuid/uuid"
	"github.com/coreos/go-etcd/etcd"

	"github.com/wantedly/risu/schema"
)

const (
	DefaultMachines  = "http://172.17.8.101:4001"
	DefaultKeyPrefix = "/risu/"
)

type EtcdRegistry struct {
	etcd      etcd.Client
	keyPrefix string
}

func NewEtcdRegistry(machines string) Registry {
	if os.Getenv("RISU_ETCD_MACHINES") != "" {
		machines = os.Getenv("RISU_ETCD_MACHINES")
	}

	if machines == "" {
		machines = DefaultMachines
	}

	m := strings.Split(machines, ",")
	etcdClient := *etcd.NewClient(m)
	return &EtcdRegistry{etcdClient, DefaultKeyPrefix}
}

func (r *EtcdRegistry) Set(build schema.Build) error {
	j, err := marshal(build)
	if err != nil {
		return err
	}

	key := path.Join(r.keyPrefix, build.ID.String())
	_, err = r.etcd.Set(key, string(j), 0)
	if err != nil {
		return err
	}

	return nil
}

func (r *EtcdRegistry) Get(id uuid.UUID) (schema.Build, error) {
	key := path.Join(r.keyPrefix, id.String())
	res, err := r.etcd.Get(key, false, true)
	if err != nil {
		return schema.Build{}, err
	}

	var build schema.Build
	err = unmarshal(res.Node.Value, &build)
	if err != nil {
		if isKeyNotFound(err) {
			// TODO (spesnova): return 404 error
			return schema.Build{}, err
		}
		return schema.Build{}, err
	}

	return build, nil
}

func marshal(obj interface{}) (string, error) {
	encoded, err := json.Marshal(obj)
	if err != nil {
		return "", fmt.Errorf("unable to JSON-serialize object: %s", err)
	}
	return string(encoded), nil
}

func unmarshal(val string, obj interface{}) error {
	err := json.Unmarshal([]byte(val), &obj)
	if err != nil {
		return fmt.Errorf("unable to JSON-deserialize object: %s", err)
	}
	return nil
}

func isKeyNotFound(err error) bool {
	e, ok := err.(*etcd.EtcdError)
	return ok && e.ErrorCode == 100
}
<commit_msg>EtcdRegistry has list method<commit_after>package registry

import (
	"encoding/json"
	"fmt"
	"os"
	"path"
	"strings"

	"code.google.com/p/go-uuid/uuid"
	"github.com/coreos/go-etcd/etcd"

	"github.com/wantedly/risu/schema"
)

const (
	DefaultMachines  = "http://172.17.8.101:4001"
	DefaultKeyPrefix = "/risu/"
)

type EtcdRegistry struct {
	etcd      etcd.Client
	keyPrefix string
}

func NewEtcdRegistry(machines string) Registry {
	if os.Getenv("RISU_ETCD_MACHINES") != "" {
		machines = os.Getenv("RISU_ETCD_MACHINES")
	}

	if machines == "" {
		machines = DefaultMachines
	}

	m := strings.Split(machines, ",")
	etcdClient := *etcd.NewClient(m)
	return &EtcdRegistry{etcdClient, DefaultKeyPrefix}
}

func (r *EtcdRegistry) Set(build schema.Build) error {
	j, err := marshal(build)
	if err != nil {
		return err
	}

	key := path.Join(r.keyPrefix, build.ID.String())
	_, err = r.etcd.Set(key, string(j), 0)
	if err != nil {
		return err
	}

	return nil
}

func (r *EtcdRegistry) Get(id uuid.UUID) (schema.Build, error) {
	key := path.Join(r.keyPrefix, id.String())
	res, err := r.etcd.Get(key, false, true)
	if err != nil {
		return schema.Build{}, err
	}

	var build schema.Build
	err = unmarshal(res.Node.Value, &build)
	if err != nil {
		if isKeyNotFound(err) {
			// TODO (spesnova): return 404 error
			return schema.Build{}, err
		}
		return schema.Build{}, err
	}

	return build, nil
}

func (r *EtcdRegistry) List() ([]schema.Build, error) {
	var builds []schema.Build

	key := path.Join(r.keyPrefix)
	res, err := r.etcd.Get(key, false, true)
	if err != nil {
		if isKeyNotFound(err) {
			return nil, nil
		}
		return nil, err
	}

	for _, node := range res.Node.Nodes {
		var build schema.Build
		err = unmarshal(node.Value, &build)
		if err != nil {
			return nil, err
		}

		builds = append(builds, build)
	}

	return builds, nil
}

func marshal(obj interface{}) (string, error) {
	encoded, err := json.Marshal(obj)
	if err != nil {
		return "", fmt.Errorf("unable to JSON-serialize object: %s", err)
	}
	return string(encoded), nil
}

func unmarshal(val string, obj interface{}) error {
	err := json.Unmarshal([]byte(val), &obj)
	if err != nil {
		return fmt.Errorf("unable to JSON-deserialize object: %s", err)
	}
	return nil
}

func isKeyNotFound(err error) bool {
	e, ok := err.(*etcd.EtcdError)
	return ok && e.ErrorCode == 100
}
<|endoftext|>


<commit_before>// utilities to work on raw message
package message

import (
	"bytes"
	"net/textproto"
	"strings"
)

// RawGetHeaders return raw headers
func RawGetHeaders(raw *[]byte) []byte {
	return bytes.Split(*raw, []byte{13, 10, 13, 10})[0]
}

// RawHaveHeader check igf header header is present in raw mail
func RawHaveHeader(raw *[]byte, header string) bool {
	var bHeader []byte
	if strings.ToLower(header) == "message-id" {
		bHeader = []byte("Message-ID")
	} else {
		bHeader = []byte(textproto.CanonicalMIMEHeaderKey(header))
	}
	for _, line := range bytes.Split(RawGetHeaders(raw), []byte{13, 10}) {
		if bytes.HasPrefix(line, bHeader) {
			return true
		}
	}
	return false
}

// RawGetMessageId return Message-ID or empty string if to found
func RawGetMessageId(raw *[]byte) []byte {
	bHeader := []byte("Message-ID")
	for _, line := range bytes.Split(RawGetHeaders(raw), []byte{13, 10}) {
		if bytes.HasPrefix(line, bHeader) {
			// strip <>
			return bytes.TrimPrefix(bytes.TrimSuffix(line[12:], []byte{62}), []byte{60})
		}
	}
	return []byte{}
}
<commit_msg>Message-ID vs Message-Id<commit_after>// utilities to work on raw message
package message

import (
	"bytes"
	"net/textproto"
	"strings"
)

// RawGetHeaders return raw headers
func RawGetHeaders(raw *[]byte) []byte {
	return bytes.Split(*raw, []byte{13, 10, 13, 10})[0]
}

// RawHaveHeader check igf header header is present in raw mail
func RawHaveHeader(raw *[]byte, header string) bool {
	var bHeader []byte
	if strings.ToLower(header) == "message-id" {
		bHeader = []byte("Message-ID")
	} else {
		bHeader = []byte(textproto.CanonicalMIMEHeaderKey(header))
	}
	for _, line := range bytes.Split(RawGetHeaders(raw), []byte{13, 10}) {
		if bytes.HasPrefix(line, bHeader) {
			return true
		}
	}
	return false
}

// RawGetMessageId return Message-ID or empty string if to found
func RawGetMessageId(raw *[]byte) []byte {
	bHeader := []byte("message-id")
	for _, line := range bytes.Split(RawGetHeaders(raw), []byte{13, 10}) {
		if bytes.HasPrefix(bytes.ToLower(line), bHeader) {
			// strip <>
			return bytes.TrimPrefix(bytes.TrimSuffix(line[12:], []byte{62}), []byte{60})
		}
	}
	return []byte{}
}
<|endoftext|>


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

// Provide thin wrappers around functions and constants given by the liblzma
// library. The liblzma library by Lasse Collin is used because there are no
// available pure Go implementations that support the newer XZ file format.
//
// This wrapper library intended to be used by the lzma and xz libraries that
// exist in the same parent package. The author of this package does not ensure
// that function names, constants names, and/or function implementation remain
// stable. Thus, use this library at your own peril!
package lib

/*
#cgo LDFLAGS: -llzma
#include "lzma.h"
*/
import "C"

import "math"

// Status code constants:
const (
	OK                = C.LZMA_OK
	STREAM_END        = C.LZMA_STREAM_END
	NO_CHECK          = C.LZMA_NO_CHECK
	UNSUPPORTED_CHECK = C.LZMA_UNSUPPORTED_CHECK
	GET_CHECK         = C.LZMA_GET_CHECK
	MEM_ERROR         = C.LZMA_MEM_ERROR
	MEMLIMIT_ERROR    = C.LZMA_MEMLIMIT_ERROR
	FORMAT_ERROR      = C.LZMA_FORMAT_ERROR
	OPTIONS_ERROR     = C.LZMA_OPTIONS_ERROR
	DATA_ERROR        = C.LZMA_DATA_ERROR
	BUF_ERROR         = C.LZMA_BUF_ERROR
	PROG_ERROR        = C.LZMA_PROG_ERROR
)

// Flushing mode constants:
const (
	RUN        = C.LZMA_RUN
	SYNC_FLUSH = C.LZMA_SYNC_FLUSH
	FULL_FLUSH = C.LZMA_FULL_FLUSH
	FINISH     = C.LZMA_FINISH
)

// Encoder flag constants:
const (
	PRESET_DEFAULT    = C.LZMA_PRESET_DEFAULT
	PRESET_LEVEL_MASK = C.LZMA_PRESET_LEVEL_MASK
	PRESET_EXTREME    = C.LZMA_PRESET_EXTREME
)

const (
	PRESET_LEVEL0 = iota
	PRESET_LEVEL1
	PRESET_LEVEL2
	PRESET_LEVEL3
	PRESET_LEVEL4
	PRESET_LEVEL5
	PRESET_LEVEL6
	PRESET_LEVEL7
	PRESET_LEVEL8
	PRESET_LEVEL9
)

// Decoder flag constants:
const (
	TELL_NO_CHECK          = C.LZMA_TELL_NO_CHECK
	TELL_UNSUPPORTED_CHECK = C.LZMA_TELL_UNSUPPORTED_CHECK
	TELL_ANY_CHECK         = C.LZMA_TELL_ANY_CHECK
	CONCATENATED           = C.LZMA_CONCATENATED
)

// Integrity check constants:
const (
	CHECK_NONE   = C.LZMA_CHECK_NONE
	CHECK_CRC32  = C.LZMA_CHECK_CRC32
	CHECK_CRC64  = C.LZMA_CHECK_CRC64
	CHECK_SHA256 = C.LZMA_CHECK_SHA256
)

// Index iteration mode constants:
const (
	INDEX_ITER_ANY            = C.LZMA_INDEX_ITER_ANY
	INDEX_ITER_BLOCK          = C.LZMA_INDEX_ITER_BLOCK
	INDEX_ITER_NONEMPTY_BLOCK = C.LZMA_INDEX_ITER_NONEMPTY_BLOCK
	INDEX_ITER_STREAM         = C.LZMA_INDEX_ITER_STREAM
)

// Variable length integer constants:
const (
	VLI_MAX       = C.LZMA_VLI_MAX
	VLI_BYTES_MAX = C.LZMA_VLI_BYTES_MAX
	VLI_UNKNOWN   = math.MaxUint64
)

// HACK(jtsai): LZMA_VLI_UNKNOWN is effectively -1 which overflows when Go tries
//  to cast the constant to the lzma_vli type which is unsigned. Since the
//  lzma_vli type is a uint64_t, then the VLI_UNKNOWN constant essentially gets
//  casted as MaxUint64.

// Filter identification consts:
const (
	FILTER_LZMA1    = C.LZMA_FILTER_LZMA1
	FILTER_LZMA2    = C.LZMA_FILTER_LZMA2
	FILTER_DELTA    = C.LZMA_FILTER_DELTA
	FILTER_X86      = C.LZMA_FILTER_X86
	FILTER_POWERPC  = C.LZMA_FILTER_POWERPC
	FILTER_IA64     = C.LZMA_FILTER_IA64
	FILTER_ARM      = C.LZMA_FILTER_ARM
	FILTER_ARMTHUMB = C.LZMA_FILTER_ARMTHUMB
	FILTER_SPARC    = C.LZMA_FILTER_SPARC
)

// File format constants:
var (
	HEADER_MAGIC = []byte{0xFD, '7', 'z', 'X', 'Z', 0x00}
	FOOTER_MAGIC = []byte{'Y', 'Z'}
)

const (
	STREAM_HEADER_SIZE    = C.LZMA_STREAM_HEADER_SIZE
	STREAM_FOOTER_SIZE    = C.LZMA_STREAM_HEADER_SIZE
	BACKWARD_SIZE_MAX     = C.LZMA_BACKWARD_SIZE_MAX
	BACKWARD_SIZE_MIN     = C.LZMA_BACKWARD_SIZE_MIN
	BLOCK_HEADER_SIZE_MAX = C.LZMA_BLOCK_HEADER_SIZE_MAX
	BLOCK_HEADER_SIZE_MIN = C.LZMA_BLOCK_HEADER_SIZE_MIN
)
<commit_msg>Trivial constants change<commit_after>// Copyright 2014, Joe Tsai. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.md file.

// Provide thin wrappers around functions and constants given by the liblzma
// library. The liblzma library by Lasse Collin is used because there are no
// available pure Go implementations that support the newer XZ file format.
//
// This wrapper library intended to be used by the lzma and xz libraries that
// exist in the same parent package. The author of this package does not ensure
// that function names, constants names, and/or function implementation remain
// stable. Thus, use this library at your own peril!
package lib

/*
#cgo LDFLAGS: -llzma
#include "lzma.h"
*/
import "C"

import "math"

// Status code constants:
const (
	OK                = C.LZMA_OK
	STREAM_END        = C.LZMA_STREAM_END
	NO_CHECK          = C.LZMA_NO_CHECK
	UNSUPPORTED_CHECK = C.LZMA_UNSUPPORTED_CHECK
	GET_CHECK         = C.LZMA_GET_CHECK
	MEM_ERROR         = C.LZMA_MEM_ERROR
	MEMLIMIT_ERROR    = C.LZMA_MEMLIMIT_ERROR
	FORMAT_ERROR      = C.LZMA_FORMAT_ERROR
	OPTIONS_ERROR     = C.LZMA_OPTIONS_ERROR
	DATA_ERROR        = C.LZMA_DATA_ERROR
	BUF_ERROR         = C.LZMA_BUF_ERROR
	PROG_ERROR        = C.LZMA_PROG_ERROR
)

// Flushing mode constants:
const (
	RUN        = C.LZMA_RUN
	SYNC_FLUSH = C.LZMA_SYNC_FLUSH
	FULL_FLUSH = C.LZMA_FULL_FLUSH
	FINISH     = C.LZMA_FINISH
)

// Encoder flag constants:
const (
	PRESET_LEVEL0 = iota
	PRESET_LEVEL1
	PRESET_LEVEL2
	PRESET_LEVEL3
	PRESET_LEVEL4
	PRESET_LEVEL5
	PRESET_LEVEL6
	PRESET_LEVEL7
	PRESET_LEVEL8
	PRESET_LEVEL9
	PRESET_DEFAULT    = C.LZMA_PRESET_DEFAULT
	PRESET_LEVEL_MASK = C.LZMA_PRESET_LEVEL_MASK
	PRESET_EXTREME    = C.LZMA_PRESET_EXTREME
)

// Decoder flag constants:
const (
	TELL_NO_CHECK          = C.LZMA_TELL_NO_CHECK
	TELL_UNSUPPORTED_CHECK = C.LZMA_TELL_UNSUPPORTED_CHECK
	TELL_ANY_CHECK         = C.LZMA_TELL_ANY_CHECK
	CONCATENATED           = C.LZMA_CONCATENATED
)

// Integrity check constants:
const (
	CHECK_NONE   = C.LZMA_CHECK_NONE
	CHECK_CRC32  = C.LZMA_CHECK_CRC32
	CHECK_CRC64  = C.LZMA_CHECK_CRC64
	CHECK_SHA256 = C.LZMA_CHECK_SHA256
)

// Index iteration mode constants:
const (
	INDEX_ITER_ANY            = C.LZMA_INDEX_ITER_ANY
	INDEX_ITER_BLOCK          = C.LZMA_INDEX_ITER_BLOCK
	INDEX_ITER_NONEMPTY_BLOCK = C.LZMA_INDEX_ITER_NONEMPTY_BLOCK
	INDEX_ITER_STREAM         = C.LZMA_INDEX_ITER_STREAM
)

// Variable length integer constants:
const (
	VLI_MAX       = C.LZMA_VLI_MAX
	VLI_BYTES_MAX = C.LZMA_VLI_BYTES_MAX
	VLI_UNKNOWN   = math.MaxUint64
)

// HACK(jtsai): LZMA_VLI_UNKNOWN is effectively -1 which overflows when Go tries
//  to cast the constant to the lzma_vli type which is unsigned. Since the
//  lzma_vli type is a uint64_t, then the VLI_UNKNOWN constant essentially gets
//  casted as MaxUint64.

// Filter identification consts:
const (
	FILTER_LZMA1    = C.LZMA_FILTER_LZMA1
	FILTER_LZMA2    = C.LZMA_FILTER_LZMA2
	FILTER_DELTA    = C.LZMA_FILTER_DELTA
	FILTER_X86      = C.LZMA_FILTER_X86
	FILTER_POWERPC  = C.LZMA_FILTER_POWERPC
	FILTER_IA64     = C.LZMA_FILTER_IA64
	FILTER_ARM      = C.LZMA_FILTER_ARM
	FILTER_ARMTHUMB = C.LZMA_FILTER_ARMTHUMB
	FILTER_SPARC    = C.LZMA_FILTER_SPARC
)

// File format constants:
var (
	HEADER_MAGIC = []byte{0xFD, '7', 'z', 'X', 'Z', 0x00}
	FOOTER_MAGIC = []byte{'Y', 'Z'}
)

const (
	STREAM_HEADER_SIZE    = C.LZMA_STREAM_HEADER_SIZE
	STREAM_FOOTER_SIZE    = C.LZMA_STREAM_HEADER_SIZE
	BACKWARD_SIZE_MAX     = C.LZMA_BACKWARD_SIZE_MAX
	BACKWARD_SIZE_MIN     = C.LZMA_BACKWARD_SIZE_MIN
	BLOCK_HEADER_SIZE_MAX = C.LZMA_BLOCK_HEADER_SIZE_MAX
	BLOCK_HEADER_SIZE_MIN = C.LZMA_BLOCK_HEADER_SIZE_MIN
)
<|endoftext|>


<commit_before>package polyline

import (
	"math"
)

func round(x float64) float64 {
	if x < 0 {
		return -math.Floor(-x + 0.5)
	} else {
		return math.Floor(x + 0.5)
	}
}

func EncodeUint(u uint, result []byte) []byte {
	for u >= 32 {
		result = append(result, byte((u&31)+95))
		u = u >> 5
	}
	result = append(result, byte(u+63))
	return result
}

func EncodeInt(i int, result []byte) []byte {
	var u uint
	if i < 0 {
		u = uint(^(i << 1))
	} else {
		u = uint(i << 1)
	}
	return EncodeUint(u, result)
}

func EncodeFloat64(f float64, result []byte) []byte {
	return EncodeInt(int(round(1e5*f)), result)
}

func EncodeCoords(coords [][]float64, result []byte) []byte {
	lastLat, lastLong := 0, 0
	for _, coord := range coords {
		lat, long := int(1e5*coord[0]), int(1e5*coord[1])
		if lat != lastLat && long != lastLong {
			result = EncodeInt(lat-lastLat, result)
			result = EncodeInt(long-lastLong, result)
			lastLat, lastLong = lat, long
		}
	}
	return result
}
<commit_msg>Add some documentation<commit_after>// Package polyline implements a Google Maps Encoding Polyline encoder.
//
// See https://developers.google.com/maps/documentation/utilities/polylinealgorithm
package polyline

import (
	"math"
)

func round(x float64) float64 {
	if x < 0 {
		return -math.Floor(-x + 0.5)
	} else {
		return math.Floor(x + 0.5)
	}
}

// EncodeUint appends the encoding of a single unsigned integer u to result.
func EncodeUint(u uint, result []byte) []byte {
	for u >= 32 {
		result = append(result, byte((u&31)+95))
		u = u >> 5
	}
	result = append(result, byte(u+63))
	return result
}

// EncodeInt appends the encoding of a single signed integer i to result.
func EncodeInt(i int, result []byte) []byte {
	var u uint
	if i < 0 {
		u = uint(^(i << 1))
	} else {
		u = uint(i << 1)
	}
	return EncodeUint(u, result)
}

// EncodeFloat64 appends the encoding of a single float64 f to result.
func EncodeFloat64(f float64, result []byte) []byte {
	return EncodeInt(int(round(1e5*f)), result)
}

// EncodeCoords appends the encoding of an array of coordinates coords to result.
func EncodeCoords(coords [][]float64, result []byte) []byte {
	lastLat, lastLong := 0, 0
	for _, coord := range coords {
		lat, long := int(1e5*coord[0]), int(1e5*coord[1])
		if lat != lastLat && long != lastLong {
			result = EncodeInt(lat-lastLat, result)
			result = EncodeInt(long-lastLong, result)
			lastLat, lastLong = lat, long
		}
	}
	return result
}
<|endoftext|>


<commit_before>// Copyright 2010 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 gzip

import (
	"bufio"
	"bytes"
	oldgz "compress/gzip"
	"io"
	"io/ioutil"
	"math/rand"
	"testing"
	"time"
)

// TestEmpty tests that an empty payload still forms a valid GZIP stream.
func TestEmpty(t *testing.T) {
	buf := new(bytes.Buffer)

	if err := NewWriter(buf).Close(); err != nil {
		t.Fatalf("Writer.Close: %v", err)
	}

	r, err := NewReader(buf)
	if err != nil {
		t.Fatalf("NewReader: %v", err)
	}
	b, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatalf("ReadAll: %v", err)
	}
	if len(b) != 0 {
		t.Fatalf("got %d bytes, want 0", len(b))
	}
	if err := r.Close(); err != nil {
		t.Fatalf("Reader.Close: %v", err)
	}
}

// TestRoundTrip tests that gzipping and then gunzipping is the identity
// function.
func TestRoundTrip(t *testing.T) {
	buf := new(bytes.Buffer)

	w := NewWriter(buf)
	w.Comment = "comment"
	w.Extra = []byte("extra")
	w.ModTime = time.Unix(1e8, 0)
	w.Name = "name"
	if _, err := w.Write([]byte("payload")); err != nil {
		t.Fatalf("Write: %v", err)
	}
	if err := w.Close(); err != nil {
		t.Fatalf("Writer.Close: %v", err)
	}

	r, err := NewReader(buf)
	if err != nil {
		t.Fatalf("NewReader: %v", err)
	}
	b, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatalf("ReadAll: %v", err)
	}
	if string(b) != "payload" {
		t.Fatalf("payload is %q, want %q", string(b), "payload")
	}
	if r.Comment != "comment" {
		t.Fatalf("comment is %q, want %q", r.Comment, "comment")
	}
	if string(r.Extra) != "extra" {
		t.Fatalf("extra is %q, want %q", r.Extra, "extra")
	}
	if r.ModTime.Unix() != 1e8 {
		t.Fatalf("mtime is %d, want %d", r.ModTime.Unix(), uint32(1e8))
	}
	if r.Name != "name" {
		t.Fatalf("name is %q, want %q", r.Name, "name")
	}
	if err := r.Close(); err != nil {
		t.Fatalf("Reader.Close: %v", err)
	}
}

// TestLatin1 tests the internal functions for converting to and from Latin-1.
func TestLatin1(t *testing.T) {
	latin1 := []byte{0xc4, 'u', 0xdf, 'e', 'r', 'u', 'n', 'g', 0}
	utf8 := "Äußerung"
	z := Reader{r: bufio.NewReader(bytes.NewReader(latin1))}
	s, err := z.readString()
	if err != nil {
		t.Fatalf("readString: %v", err)
	}
	if s != utf8 {
		t.Fatalf("read latin-1: got %q, want %q", s, utf8)
	}

	buf := bytes.NewBuffer(make([]byte, 0, len(latin1)))
	c := Writer{w: buf}
	if err = c.writeString(utf8); err != nil {
		t.Fatalf("writeString: %v", err)
	}
	s = buf.String()
	if s != string(latin1) {
		t.Fatalf("write utf-8: got %q, want %q", s, string(latin1))
	}
}

// TestLatin1RoundTrip tests that metadata that is representable in Latin-1
// survives a round trip.
func TestLatin1RoundTrip(t *testing.T) {
	testCases := []struct {
		name string
		ok   bool
	}{
		{"", true},
		{"ASCII is OK", true},
		{"unless it contains a NUL\x00", false},
		{"no matter where \x00 occurs", false},
		{"\x00\x00\x00", false},
		{"Látin-1 also passes (U+00E1)", true},
		{"but LĀtin Extended-A (U+0100) does not", false},
		{"neither does 日本語", false},
		{"invalid UTF-8 also \xffails", false},
		{"\x00 as does Látin-1 with NUL", false},
	}
	for _, tc := range testCases {
		buf := new(bytes.Buffer)

		w := NewWriter(buf)
		w.Name = tc.name
		err := w.Close()
		if (err == nil) != tc.ok {
			t.Errorf("Writer.Close: name = %q, err = %v", tc.name, err)
			continue
		}
		if !tc.ok {
			continue
		}

		r, err := NewReader(buf)
		if err != nil {
			t.Errorf("NewReader: %v", err)
			continue
		}
		_, err = ioutil.ReadAll(r)
		if err != nil {
			t.Errorf("ReadAll: %v", err)
			continue
		}
		if r.Name != tc.name {
			t.Errorf("name is %q, want %q", r.Name, tc.name)
			continue
		}
		if err := r.Close(); err != nil {
			t.Errorf("Reader.Close: %v", err)
			continue
		}
	}
}

func TestWriterFlush(t *testing.T) {
	buf := new(bytes.Buffer)

	w := NewWriter(buf)
	w.Comment = "comment"
	w.Extra = []byte("extra")
	w.ModTime = time.Unix(1e8, 0)
	w.Name = "name"

	n0 := buf.Len()
	if n0 != 0 {
		t.Fatalf("buffer size = %d before writes; want 0", n0)
	}

	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}

	n1 := buf.Len()
	if n1 == 0 {
		t.Fatal("no data after first flush")
	}

	w.Write([]byte("x"))

	n2 := buf.Len()
	if n1 != n2 {
		t.Fatalf("after writing a single byte, size changed from %d to %d; want no change", n1, n2)
	}

	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}

	n3 := buf.Len()
	if n2 == n3 {
		t.Fatal("Flush didn't flush any data")
	}
}

// Multiple gzip files concatenated form a valid gzip file.
func TestConcat(t *testing.T) {
	var buf bytes.Buffer
	w := NewWriter(&buf)
	w.Write([]byte("hello "))
	w.Close()
	w = NewWriter(&buf)
	w.Write([]byte("world\n"))
	w.Close()

	r, err := NewReader(&buf)
	data, err := ioutil.ReadAll(r)
	if string(data) != "hello world\n" || err != nil {
		t.Fatalf("ReadAll = %q, %v, want %q, nil", data, err, "hello world")
	}
}

func TestWriterReset(t *testing.T) {
	buf := new(bytes.Buffer)
	buf2 := new(bytes.Buffer)
	z := NewWriter(buf)
	msg := []byte("hello world")
	z.Write(msg)
	z.Close()
	z.Reset(buf2)
	z.Write(msg)
	z.Close()
	if buf.String() != buf2.String() {
		t.Errorf("buf2 %q != original buf of %q", buf2.String(), buf.String())
	}
}

var testbuf []byte

func testFile(i int, t *testing.T) {
	dat, _ := ioutil.ReadFile("testdata/test.json")
	dl := len(dat)
	if len(testbuf) != i*dl {
		// Make results predictable
		testbuf = make([]byte, i*dl)
		for j := 0; j < i; j++ {
			copy(testbuf[j*dl:j*dl+dl], dat)
		}
	}

	br := bytes.NewBuffer(testbuf)
	var buf bytes.Buffer
	w, _ := NewWriterLevel(&buf, 6)
	io.Copy(w, br)
	w.Close()
	r, err := NewReader(&buf)
	if err != nil {
		t.Fatal(err.Error())
	}
	decoded, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatal(err.Error())
	}
	if !bytes.Equal(testbuf, decoded) {
		t.Errorf("decoded content does not match.")
	}
}

func TestFile1(t *testing.T)  { testFile(1, t) }
func TestFile10(t *testing.T) { testFile(10, t) }

func TestFile50(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}
	testFile(50, t)
}

func TestFile200(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}
	testFile(200, t)
}

func testBigGzip(i int, t *testing.T) {
	if len(testbuf) != i {
		// Make results predictable
		rand.Seed(1337)
		testbuf = make([]byte, i)
		for idx := range testbuf {
			testbuf[idx] = byte(65 + rand.Intn(32))
		}
	}

	br := bytes.NewBuffer(testbuf)
	var buf bytes.Buffer
	w, _ := NewWriterLevel(&buf, 6)
	io.Copy(w, br)
	err := w.Close()
	if err != nil {
		t.Fatal(err.Error())
	}

	r, err := NewReader(&buf)
	if err != nil {
		t.Fatal(err.Error())
	}
	decoded, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatal(err.Error())
	}
	if !bytes.Equal(testbuf, decoded) {
		t.Errorf("decoded content does not match.")
	}
}

func TestGzip1K(t *testing.T)   { testBigGzip(1000, t) }
func TestGzip100K(t *testing.T) { testBigGzip(100000, t) }
func TestGzip1M(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}

	testBigGzip(1000000, t)
}
func TestGzip10M(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}
	testBigGzip(10000000, t)
}

func BenchmarkGzipL1(b *testing.B) { benchmarkGzipN(b, 1) }
func BenchmarkGzipL2(b *testing.B) { benchmarkGzipN(b, 2) }
func BenchmarkGzipL3(b *testing.B) { benchmarkGzipN(b, 3) }
func BenchmarkGzipL4(b *testing.B) { benchmarkGzipN(b, 4) }
func BenchmarkGzipL5(b *testing.B) { benchmarkGzipN(b, 5) }
func BenchmarkGzipL6(b *testing.B) { benchmarkGzipN(b, 6) }
func BenchmarkGzipL7(b *testing.B) { benchmarkGzipN(b, 7) }
func BenchmarkGzipL8(b *testing.B) { benchmarkGzipN(b, 8) }
func BenchmarkGzipL9(b *testing.B) { benchmarkGzipN(b, 9) }

func benchmarkGzipN(b *testing.B, level int) {
	dat, _ := ioutil.ReadFile("testdata/test.json")
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	b.SetBytes(int64(len(dat)))
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		w, _ := NewWriterLevel(ioutil.Discard, level)
		w.Write(dat)
		w.Flush()
		w.Close()
	}
}

func BenchmarkOldGzipL1(b *testing.B) { benchmarkOldGzipN(b, 1) }
func BenchmarkOldGzipL2(b *testing.B) { benchmarkOldGzipN(b, 2) }
func BenchmarkOldGzipL3(b *testing.B) { benchmarkOldGzipN(b, 3) }
func BenchmarkOldGzipL4(b *testing.B) { benchmarkOldGzipN(b, 4) }
func BenchmarkOldGzipL5(b *testing.B) { benchmarkOldGzipN(b, 5) }
func BenchmarkOldGzipL6(b *testing.B) { benchmarkOldGzipN(b, 6) }
func BenchmarkOldGzipL7(b *testing.B) { benchmarkOldGzipN(b, 7) }
func BenchmarkOldGzipL8(b *testing.B) { benchmarkOldGzipN(b, 8) }
func BenchmarkOldGzipL9(b *testing.B) { benchmarkOldGzipN(b, 9) }

func benchmarkOldGzipN(b *testing.B, level int) {
	dat, _ := ioutil.ReadFile("testdata/test.json")
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)

	b.SetBytes(int64(len(dat)))
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		w, _ := oldgz.NewWriterLevel(ioutil.Discard, level)
		w.Write(dat)
		w.Flush()
		w.Close()
	}
}
<commit_msg>Re-use encoder.<commit_after>// Copyright 2010 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 gzip

import (
	"bufio"
	"bytes"
	oldgz "compress/gzip"
	"io"
	"io/ioutil"
	"math/rand"
	"testing"
	"time"
)

// TestEmpty tests that an empty payload still forms a valid GZIP stream.
func TestEmpty(t *testing.T) {
	buf := new(bytes.Buffer)

	if err := NewWriter(buf).Close(); err != nil {
		t.Fatalf("Writer.Close: %v", err)
	}

	r, err := NewReader(buf)
	if err != nil {
		t.Fatalf("NewReader: %v", err)
	}
	b, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatalf("ReadAll: %v", err)
	}
	if len(b) != 0 {
		t.Fatalf("got %d bytes, want 0", len(b))
	}
	if err := r.Close(); err != nil {
		t.Fatalf("Reader.Close: %v", err)
	}
}

// TestRoundTrip tests that gzipping and then gunzipping is the identity
// function.
func TestRoundTrip(t *testing.T) {
	buf := new(bytes.Buffer)

	w := NewWriter(buf)
	w.Comment = "comment"
	w.Extra = []byte("extra")
	w.ModTime = time.Unix(1e8, 0)
	w.Name = "name"
	if _, err := w.Write([]byte("payload")); err != nil {
		t.Fatalf("Write: %v", err)
	}
	if err := w.Close(); err != nil {
		t.Fatalf("Writer.Close: %v", err)
	}

	r, err := NewReader(buf)
	if err != nil {
		t.Fatalf("NewReader: %v", err)
	}
	b, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatalf("ReadAll: %v", err)
	}
	if string(b) != "payload" {
		t.Fatalf("payload is %q, want %q", string(b), "payload")
	}
	if r.Comment != "comment" {
		t.Fatalf("comment is %q, want %q", r.Comment, "comment")
	}
	if string(r.Extra) != "extra" {
		t.Fatalf("extra is %q, want %q", r.Extra, "extra")
	}
	if r.ModTime.Unix() != 1e8 {
		t.Fatalf("mtime is %d, want %d", r.ModTime.Unix(), uint32(1e8))
	}
	if r.Name != "name" {
		t.Fatalf("name is %q, want %q", r.Name, "name")
	}
	if err := r.Close(); err != nil {
		t.Fatalf("Reader.Close: %v", err)
	}
}

// TestLatin1 tests the internal functions for converting to and from Latin-1.
func TestLatin1(t *testing.T) {
	latin1 := []byte{0xc4, 'u', 0xdf, 'e', 'r', 'u', 'n', 'g', 0}
	utf8 := "Äußerung"
	z := Reader{r: bufio.NewReader(bytes.NewReader(latin1))}
	s, err := z.readString()
	if err != nil {
		t.Fatalf("readString: %v", err)
	}
	if s != utf8 {
		t.Fatalf("read latin-1: got %q, want %q", s, utf8)
	}

	buf := bytes.NewBuffer(make([]byte, 0, len(latin1)))
	c := Writer{w: buf}
	if err = c.writeString(utf8); err != nil {
		t.Fatalf("writeString: %v", err)
	}
	s = buf.String()
	if s != string(latin1) {
		t.Fatalf("write utf-8: got %q, want %q", s, string(latin1))
	}
}

// TestLatin1RoundTrip tests that metadata that is representable in Latin-1
// survives a round trip.
func TestLatin1RoundTrip(t *testing.T) {
	testCases := []struct {
		name string
		ok   bool
	}{
		{"", true},
		{"ASCII is OK", true},
		{"unless it contains a NUL\x00", false},
		{"no matter where \x00 occurs", false},
		{"\x00\x00\x00", false},
		{"Látin-1 also passes (U+00E1)", true},
		{"but LĀtin Extended-A (U+0100) does not", false},
		{"neither does 日本語", false},
		{"invalid UTF-8 also \xffails", false},
		{"\x00 as does Látin-1 with NUL", false},
	}
	for _, tc := range testCases {
		buf := new(bytes.Buffer)

		w := NewWriter(buf)
		w.Name = tc.name
		err := w.Close()
		if (err == nil) != tc.ok {
			t.Errorf("Writer.Close: name = %q, err = %v", tc.name, err)
			continue
		}
		if !tc.ok {
			continue
		}

		r, err := NewReader(buf)
		if err != nil {
			t.Errorf("NewReader: %v", err)
			continue
		}
		_, err = ioutil.ReadAll(r)
		if err != nil {
			t.Errorf("ReadAll: %v", err)
			continue
		}
		if r.Name != tc.name {
			t.Errorf("name is %q, want %q", r.Name, tc.name)
			continue
		}
		if err := r.Close(); err != nil {
			t.Errorf("Reader.Close: %v", err)
			continue
		}
	}
}

func TestWriterFlush(t *testing.T) {
	buf := new(bytes.Buffer)

	w := NewWriter(buf)
	w.Comment = "comment"
	w.Extra = []byte("extra")
	w.ModTime = time.Unix(1e8, 0)
	w.Name = "name"

	n0 := buf.Len()
	if n0 != 0 {
		t.Fatalf("buffer size = %d before writes; want 0", n0)
	}

	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}

	n1 := buf.Len()
	if n1 == 0 {
		t.Fatal("no data after first flush")
	}

	w.Write([]byte("x"))

	n2 := buf.Len()
	if n1 != n2 {
		t.Fatalf("after writing a single byte, size changed from %d to %d; want no change", n1, n2)
	}

	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}

	n3 := buf.Len()
	if n2 == n3 {
		t.Fatal("Flush didn't flush any data")
	}
}

// Multiple gzip files concatenated form a valid gzip file.
func TestConcat(t *testing.T) {
	var buf bytes.Buffer
	w := NewWriter(&buf)
	w.Write([]byte("hello "))
	w.Close()
	w = NewWriter(&buf)
	w.Write([]byte("world\n"))
	w.Close()

	r, err := NewReader(&buf)
	data, err := ioutil.ReadAll(r)
	if string(data) != "hello world\n" || err != nil {
		t.Fatalf("ReadAll = %q, %v, want %q, nil", data, err, "hello world")
	}
}

func TestWriterReset(t *testing.T) {
	buf := new(bytes.Buffer)
	buf2 := new(bytes.Buffer)
	z := NewWriter(buf)
	msg := []byte("hello world")
	z.Write(msg)
	z.Close()
	z.Reset(buf2)
	z.Write(msg)
	z.Close()
	if buf.String() != buf2.String() {
		t.Errorf("buf2 %q != original buf of %q", buf2.String(), buf.String())
	}
}

var testbuf []byte

func testFile(i int, t *testing.T) {
	dat, _ := ioutil.ReadFile("testdata/test.json")
	dl := len(dat)
	if len(testbuf) != i*dl {
		// Make results predictable
		testbuf = make([]byte, i*dl)
		for j := 0; j < i; j++ {
			copy(testbuf[j*dl:j*dl+dl], dat)
		}
	}

	br := bytes.NewBuffer(testbuf)
	var buf bytes.Buffer
	w, _ := NewWriterLevel(&buf, 6)
	io.Copy(w, br)
	w.Close()
	r, err := NewReader(&buf)
	if err != nil {
		t.Fatal(err.Error())
	}
	decoded, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatal(err.Error())
	}
	if !bytes.Equal(testbuf, decoded) {
		t.Errorf("decoded content does not match.")
	}
}

func TestFile1(t *testing.T)  { testFile(1, t) }
func TestFile10(t *testing.T) { testFile(10, t) }

func TestFile50(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}
	testFile(50, t)
}

func TestFile200(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}
	testFile(200, t)
}

func testBigGzip(i int, t *testing.T) {
	if len(testbuf) != i {
		// Make results predictable
		rand.Seed(1337)
		testbuf = make([]byte, i)
		for idx := range testbuf {
			testbuf[idx] = byte(65 + rand.Intn(32))
		}
	}

	br := bytes.NewBuffer(testbuf)
	var buf bytes.Buffer
	w, _ := NewWriterLevel(&buf, 6)
	io.Copy(w, br)
	err := w.Close()
	if err != nil {
		t.Fatal(err.Error())
	}

	r, err := NewReader(&buf)
	if err != nil {
		t.Fatal(err.Error())
	}
	decoded, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatal(err.Error())
	}
	if !bytes.Equal(testbuf, decoded) {
		t.Errorf("decoded content does not match.")
	}
}

func TestGzip1K(t *testing.T)   { testBigGzip(1000, t) }
func TestGzip100K(t *testing.T) { testBigGzip(100000, t) }
func TestGzip1M(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}

	testBigGzip(1000000, t)
}
func TestGzip10M(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}
	testBigGzip(10000000, t)
}

func BenchmarkGzipL1(b *testing.B) { benchmarkGzipN(b, 1) }
func BenchmarkGzipL2(b *testing.B) { benchmarkGzipN(b, 2) }
func BenchmarkGzipL3(b *testing.B) { benchmarkGzipN(b, 3) }
func BenchmarkGzipL4(b *testing.B) { benchmarkGzipN(b, 4) }
func BenchmarkGzipL5(b *testing.B) { benchmarkGzipN(b, 5) }
func BenchmarkGzipL6(b *testing.B) { benchmarkGzipN(b, 6) }
func BenchmarkGzipL7(b *testing.B) { benchmarkGzipN(b, 7) }
func BenchmarkGzipL8(b *testing.B) { benchmarkGzipN(b, 8) }
func BenchmarkGzipL9(b *testing.B) { benchmarkGzipN(b, 9) }

func benchmarkGzipN(b *testing.B, level int) {
	dat, _ := ioutil.ReadFile("testdata/test.json")
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	b.SetBytes(int64(len(dat)))
	w, _ := NewWriterLevel(ioutil.Discard, level)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		w.Reset(ioutil.Discard)
		w.Write(dat)
		w.Flush()
		w.Close()
	}
}

func BenchmarkOldGzipL1(b *testing.B) { benchmarkOldGzipN(b, 1) }
func BenchmarkOldGzipL2(b *testing.B) { benchmarkOldGzipN(b, 2) }
func BenchmarkOldGzipL3(b *testing.B) { benchmarkOldGzipN(b, 3) }
func BenchmarkOldGzipL4(b *testing.B) { benchmarkOldGzipN(b, 4) }
func BenchmarkOldGzipL5(b *testing.B) { benchmarkOldGzipN(b, 5) }
func BenchmarkOldGzipL6(b *testing.B) { benchmarkOldGzipN(b, 6) }
func BenchmarkOldGzipL7(b *testing.B) { benchmarkOldGzipN(b, 7) }
func BenchmarkOldGzipL8(b *testing.B) { benchmarkOldGzipN(b, 8) }
func BenchmarkOldGzipL9(b *testing.B) { benchmarkOldGzipN(b, 9) }

func benchmarkOldGzipN(b *testing.B, level int) {
	dat, _ := ioutil.ReadFile("testdata/test.json")
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)
	dat = append(dat, dat...)

	b.SetBytes(int64(len(dat)))
	w, _ := oldgz.NewWriterLevel(ioutil.Discard, level)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		w.Reset(ioutil.Discard)
		w.Write(dat)
		w.Flush()
		w.Close()
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 clair 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 main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/clair/api/v1"
)

const (
	postLayerURI        = "/v1/layers"
	getLayerFeaturesURI = "/v1/layers/%s?vulnerabilities"
	httpPort            = 9279
)

func main() {
	// Parse command-line arguments.
	endpoint := flag.String("endpoint", "http://127.0.0.1:6060", "Address to Clair API")
	myAddress := flag.String("my-address", "127.0.0.1", "Address from the point of view of Clair")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] image-id\n\nOptions:\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.Parse()

	if len(flag.Args()) != 1 {
		flag.Usage()
		os.Exit(1)
	}
	imageName := flag.Args()[0]

	// Save image.
	fmt.Printf("Saving %s\n", imageName)
	path, err := save(imageName)
	defer os.RemoveAll(path)
	if err != nil {
		fmt.Printf("- Could not save image: %s\n", err)
		os.Exit(1)
	}

	// Retrieve history.
	fmt.Println("Getting image's history")
	layerIDs, err := historyFromManifest(path)
	if err != nil {
		layerIDs, err = historyFromCommand(imageName)
	}
	if err != nil || len(layerIDs) == 0 {
		fmt.Printf("- Could not get image's history: %s\n", err)
		os.Exit(1)
	}

	// Setup a simple HTTP server if Clair is not local.
	if !strings.Contains(*endpoint, "127.0.0.1") && !strings.Contains(*endpoint, "localhost") {
		allowedHost := strings.TrimPrefix(*endpoint, "http://")
		portIndex := strings.Index(allowedHost, ":")
		if portIndex >= 0 {
			allowedHost = allowedHost[:portIndex]
		}

		go listenHTTP(path, allowedHost)

		path = "http://" + *myAddress + ":" + strconv.Itoa(httpPort)
		time.Sleep(200 * time.Millisecond)
	}

	// Analyze layers.
	fmt.Printf("Analyzing %d layers\n", len(layerIDs))
	for i := 0; i < len(layerIDs); i++ {
		fmt.Printf("- Analyzing %s\n", layerIDs[i])

		var err error
		if i > 0 {
			err = analyzeLayer(*endpoint, path+"/"+layerIDs[i]+"/layer.tar", layerIDs[i], layerIDs[i-1])
		} else {
			err = analyzeLayer(*endpoint, path+"/"+layerIDs[i]+"/layer.tar", layerIDs[i], "")
		}
		if err != nil {
			fmt.Printf("- Could not analyze layer: %s\n", err)
			os.Exit(1)
		}
	}

	// Get vulnerabilities.
	fmt.Println("Getting image's vulnerabilities")
	layer, err := getLayer(*endpoint, layerIDs[len(layerIDs)-1])
	if err != nil {
		fmt.Printf("- Could not get layer information: %s\n", err)
		os.Exit(1)
	}

	// Print report.
	fmt.Printf("\n# Clair report for image %s (%s)\n", imageName, time.Now().UTC())

	if len(layer.Features) == 0 {
		fmt.Println("No feature has been detected on the image.")
		fmt.Println("This usually means that the image isn't supported by Clair.")
		os.Exit(0)
	}

	isSafe := true
	for _, feature := range layer.Features {
		fmt.Printf("## Feature: %s %s (%s)\n", feature.Name, feature.Version, feature.Namespace)

		if len(feature.Vulnerabilities) > 0 {
			isSafe = false

			fmt.Printf("   - Added by: %s\n", feature.AddedBy)

			for _, vulnerability := range feature.Vulnerabilities {
				fmt.Printf("### (%s) %s\n", vulnerability.Severity, vulnerability.Name)

				if vulnerability.Description != "" {
					fmt.Printf("    - Link:          %s\n", vulnerability.Link)
				}

				if vulnerability.Link != "" {
					fmt.Printf("    - Description:   %s\n", vulnerability.Description)
				}

				if vulnerability.FixedBy != "" {
					fmt.Printf("    - Fixed version: %s\n", vulnerability.FixedBy)
				}

				if len(vulnerability.Metadata) > 0 {
					fmt.Printf("    - Metadata:      %+v\n", vulnerability.Metadata)
				}
			}
		}
	}

	if isSafe {
		fmt.Println("\nBravo, your image looks SAFE !")
	}
}

func save(imageName string) (string, error) {
	path, err := ioutil.TempDir("", "analyze-local-image-")
	if err != nil {
		return "", err
	}

	var stderr bytes.Buffer
	save := exec.Command("docker", "save", imageName)
	save.Stderr = &stderr
	extract := exec.Command("tar", "xf", "-", "-C"+path)
	extract.Stderr = &stderr
	pipe, err := extract.StdinPipe()
	if err != nil {
		return "", err
	}
	save.Stdout = pipe

	err = extract.Start()
	if err != nil {
		return "", errors.New(stderr.String())
	}
	err = save.Run()
	if err != nil {
		return "", errors.New(stderr.String())
	}
	err = pipe.Close()
	if err != nil {
		return "", err
	}
	err = extract.Wait()
	if err != nil {
		return "", errors.New(stderr.String())
	}

	return path, nil
}

func historyFromManifest(path string) ([]string, error) {
	mf, err := os.Open(path + "/manifest.json")
	if err != nil {
		return nil, err
	}
	defer mf.Close()

	// https://github.com/docker/docker/blob/master/image/tarexport/tarexport.go#L17
	type manifestItem struct {
		Config   string
		RepoTags []string
		Layers   []string
	}

	var manifest []manifestItem
	if err = json.NewDecoder(mf).Decode(&manifest); err != nil {
		return nil, err
	} else if len(manifest) != 1 {
		return nil, err
	}
	var layers []string
	for _, layer := range manifest[0].Layers {
		layers = append(layers, strings.TrimSuffix(layer, "/layer.tar"))
	}
	return layers, nil
}

func historyFromCommand(imageName string) ([]string, error) {
	var stderr bytes.Buffer
	cmd := exec.Command("docker", "history", "-q", "--no-trunc", imageName)
	cmd.Stderr = &stderr
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return []string{}, err
	}

	err = cmd.Start()
	if err != nil {
		return []string{}, errors.New(stderr.String())
	}

	var layers []string
	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		layers = append(layers, scanner.Text())
	}

	for i := len(layers)/2 - 1; i >= 0; i-- {
		opp := len(layers) - 1 - i
		layers[i], layers[opp] = layers[opp], layers[i]
	}

	return layers, nil
}

func listenHTTP(path, allowedHost string) {
	fmt.Printf("Setting up HTTP server (allowing: %s)\n", allowedHost)

	restrictedFileServer := func(path, allowedHost string) http.Handler {
		fc := func(w http.ResponseWriter, r *http.Request) {
			if r.Host == allowedHost {
				http.FileServer(http.Dir(path)).ServeHTTP(w, r)
				return
			}
			w.WriteHeader(403)
		}
		return http.HandlerFunc(fc)
	}

	err := http.ListenAndServe(":"+strconv.Itoa(httpPort), restrictedFileServer(path, allowedHost))
	if err != nil {
		fmt.Printf("- An error occurs with the HTTP server: %s\n", err)
		os.Exit(1)
	}
}

func analyzeLayer(endpoint, path, layerName, parentLayerName string) error {
	payload := v1.LayerEnvelope{
		Layer: &v1.Layer{
			Name:       layerName,
			Path:       path,
			ParentName: parentLayerName,
			Format:     "Docker",
		},
	}

	jsonPayload, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	request, err := http.NewRequest("POST", endpoint+postLayerURI, bytes.NewBuffer(jsonPayload))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()

	if response.StatusCode != 201 {
		body, _ := ioutil.ReadAll(response.Body)
		return fmt.Errorf("- Got response %d with message %s", response.StatusCode, string(body))
	}

	return nil
}

func getLayer(endpoint, layerID string) (v1.Layer, error) {
	response, err := http.Get(endpoint + fmt.Sprintf(getLayerFeaturesURI, layerID))
	if err != nil {
		return v1.Layer{}, err
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		body, _ := ioutil.ReadAll(response.Body)
		err := fmt.Errorf("- Got response %d with message %s", response.StatusCode, string(body))
		return v1.Layer{}, err
	}

	var apiResponse v1.LayerEnvelope
	if err = json.NewDecoder(response.Body).Decode(&apiResponse); err != nil {
		return v1.Layer{}, err
	} else if apiResponse.Error != nil {
		return v1.Layer{}, errors.New(apiResponse.Error.Message)
	}

	return *apiResponse.Layer, nil
}
<commit_msg>fix allowhost<commit_after>// Copyright 2015 clair 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 main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/clair/api/v1"
)

const (
	postLayerURI        = "/v1/layers"
	getLayerFeaturesURI = "/v1/layers/%s?vulnerabilities"
	httpPort            = 9279
)

func main() {
	// Parse command-line arguments.
	endpoint := flag.String("endpoint", "http://127.0.0.1:6060", "Address to Clair API")
	myAddress := flag.String("my-address", "127.0.0.1", "Address from the point of view of Clair")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] image-id\n\nOptions:\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.Parse()

	if len(flag.Args()) != 1 {
		flag.Usage()
		os.Exit(1)
	}
	imageName := flag.Args()[0]

	// Save image.
	fmt.Printf("Saving %s\n", imageName)
	path, err := save(imageName)
	defer os.RemoveAll(path)
	if err != nil {
		fmt.Printf("- Could not save image: %s\n", err)
		os.Exit(1)
	}

	// Retrieve history.
	fmt.Println("Getting image's history")
	layerIDs, err := historyFromManifest(path)
	if err != nil {
		layerIDs, err = historyFromCommand(imageName)
	}
	if err != nil || len(layerIDs) == 0 {
		fmt.Printf("- Could not get image's history: %s\n", err)
		os.Exit(1)
	}

	// Setup a simple HTTP server if Clair is not local.
	if !strings.Contains(*endpoint, "127.0.0.1") && !strings.Contains(*endpoint, "localhost") {
		allowedHost := strings.TrimPrefix(*endpoint, "http://")
		portIndex := strings.Index(allowedHost, ":")
		if portIndex >= 0 {
			allowedHost = allowedHost[:portIndex]
		}

		go listenHTTP(path, allowedHost)

		path = "http://" + *myAddress + ":" + strconv.Itoa(httpPort)
		time.Sleep(200 * time.Millisecond)
	}

	// Analyze layers.
	fmt.Printf("Analyzing %d layers\n", len(layerIDs))
	for i := 0; i < len(layerIDs); i++ {
		fmt.Printf("- Analyzing %s\n", layerIDs[i])

		var err error
		if i > 0 {
			err = analyzeLayer(*endpoint, path+"/"+layerIDs[i]+"/layer.tar", layerIDs[i], layerIDs[i-1])
		} else {
			err = analyzeLayer(*endpoint, path+"/"+layerIDs[i]+"/layer.tar", layerIDs[i], "")
		}
		if err != nil {
			fmt.Printf("- Could not analyze layer: %s\n", err)
			os.Exit(1)
		}
	}

	// Get vulnerabilities.
	fmt.Println("Getting image's vulnerabilities")
	layer, err := getLayer(*endpoint, layerIDs[len(layerIDs)-1])
	if err != nil {
		fmt.Printf("- Could not get layer information: %s\n", err)
		os.Exit(1)
	}

	// Print report.
	fmt.Printf("\n# Clair report for image %s (%s)\n", imageName, time.Now().UTC())

	if len(layer.Features) == 0 {
		fmt.Println("No feature has been detected on the image.")
		fmt.Println("This usually means that the image isn't supported by Clair.")
		os.Exit(0)
	}

	isSafe := true
	for _, feature := range layer.Features {
		fmt.Printf("## Feature: %s %s (%s)\n", feature.Name, feature.Version, feature.Namespace)

		if len(feature.Vulnerabilities) > 0 {
			isSafe = false

			fmt.Printf("   - Added by: %s\n", feature.AddedBy)

			for _, vulnerability := range feature.Vulnerabilities {
				fmt.Printf("### (%s) %s\n", vulnerability.Severity, vulnerability.Name)

				if vulnerability.Description != "" {
					fmt.Printf("    - Link:          %s\n", vulnerability.Link)
				}

				if vulnerability.Link != "" {
					fmt.Printf("    - Description:   %s\n", vulnerability.Description)
				}

				if vulnerability.FixedBy != "" {
					fmt.Printf("    - Fixed version: %s\n", vulnerability.FixedBy)
				}

				if len(vulnerability.Metadata) > 0 {
					fmt.Printf("    - Metadata:      %+v\n", vulnerability.Metadata)
				}
			}
		}
	}

	if isSafe {
		fmt.Println("\nBravo, your image looks SAFE !")
	}
}

func save(imageName string) (string, error) {
	path, err := ioutil.TempDir("", "analyze-local-image-")
	if err != nil {
		return "", err
	}

	var stderr bytes.Buffer
	save := exec.Command("docker", "save", imageName)
	save.Stderr = &stderr
	extract := exec.Command("tar", "xf", "-", "-C"+path)
	extract.Stderr = &stderr
	pipe, err := extract.StdinPipe()
	if err != nil {
		return "", err
	}
	save.Stdout = pipe

	err = extract.Start()
	if err != nil {
		return "", errors.New(stderr.String())
	}
	err = save.Run()
	if err != nil {
		return "", errors.New(stderr.String())
	}
	err = pipe.Close()
	if err != nil {
		return "", err
	}
	err = extract.Wait()
	if err != nil {
		return "", errors.New(stderr.String())
	}

	return path, nil
}

func historyFromManifest(path string) ([]string, error) {
	mf, err := os.Open(path + "/manifest.json")
	if err != nil {
		return nil, err
	}
	defer mf.Close()

	// https://github.com/docker/docker/blob/master/image/tarexport/tarexport.go#L17
	type manifestItem struct {
		Config   string
		RepoTags []string
		Layers   []string
	}

	var manifest []manifestItem
	if err = json.NewDecoder(mf).Decode(&manifest); err != nil {
		return nil, err
	} else if len(manifest) != 1 {
		return nil, err
	}
	var layers []string
	for _, layer := range manifest[0].Layers {
		layers = append(layers, strings.TrimSuffix(layer, "/layer.tar"))
	}
	return layers, nil
}

func historyFromCommand(imageName string) ([]string, error) {
	var stderr bytes.Buffer
	cmd := exec.Command("docker", "history", "-q", "--no-trunc", imageName)
	cmd.Stderr = &stderr
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return []string{}, err
	}

	err = cmd.Start()
	if err != nil {
		return []string{}, errors.New(stderr.String())
	}

	var layers []string
	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		layers = append(layers, scanner.Text())
	}

	for i := len(layers)/2 - 1; i >= 0; i-- {
		opp := len(layers) - 1 - i
		layers[i], layers[opp] = layers[opp], layers[i]
	}

	return layers, nil
}

func listenHTTP(path, allowedHost string) {
	fmt.Printf("Setting up HTTP server (allowing: %s)\n", allowedHost)

	restrictedFileServer := func(path, allowedHost string) http.Handler {
		fc := func(w http.ResponseWriter, r *http.Request) {
			host, _, err := net.SplitHostPort(r.RemoteAddr)
			if err == nil && strings.EqualFold(host, allowedHost) {
				http.FileServer(http.Dir(path)).ServeHTTP(w, r)
				return
			}
			w.WriteHeader(403)
		}
		return http.HandlerFunc(fc)
	}

	err := http.ListenAndServe(":"+strconv.Itoa(httpPort), restrictedFileServer(path, allowedHost))
	if err != nil {
		fmt.Printf("- An error occurs with the HTTP server: %s\n", err)
		os.Exit(1)
	}
}

func analyzeLayer(endpoint, path, layerName, parentLayerName string) error {
	payload := v1.LayerEnvelope{
		Layer: &v1.Layer{
			Name:       layerName,
			Path:       path,
			ParentName: parentLayerName,
			Format:     "Docker",
		},
	}

	jsonPayload, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	request, err := http.NewRequest("POST", endpoint+postLayerURI, bytes.NewBuffer(jsonPayload))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()

	if response.StatusCode != 201 {
		body, _ := ioutil.ReadAll(response.Body)
		return fmt.Errorf("- Got response %d with message %s", response.StatusCode, string(body))
	}

	return nil
}

func getLayer(endpoint, layerID string) (v1.Layer, error) {
	response, err := http.Get(endpoint + fmt.Sprintf(getLayerFeaturesURI, layerID))
	if err != nil {
		return v1.Layer{}, err
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		body, _ := ioutil.ReadAll(response.Body)
		err := fmt.Errorf("- Got response %d with message %s", response.StatusCode, string(body))
		return v1.Layer{}, err
	}

	var apiResponse v1.LayerEnvelope
	if err = json.NewDecoder(response.Body).Decode(&apiResponse); err != nil {
		return v1.Layer{}, err
	} else if apiResponse.Error != nil {
		return v1.Layer{}, errors.New(apiResponse.Error.Message)
	}

	return *apiResponse.Layer, nil
}
<|endoftext|>


<commit_before>package client

import (
	"errors"
	"fmt"
	"os"

	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pps"

	google_protobuf "go.pedge.io/pb/go/google/protobuf"
)

// PfsAPIClient is an alias for pfs.APIClient.
type PfsAPIClient pfs.APIClient

// PpsAPIClient is an alias for pps.APIClient.
type PpsAPIClient pps.APIClient

// BlockAPIClient is an alias for pfs.BlockAPIClient.
type BlockAPIClient pfs.BlockAPIClient

// An APIClient is a wrapper around pfs, pps and block APIClients.
type APIClient struct {
	PfsAPIClient
	PpsAPIClient
	BlockAPIClient
}

// NewFromAddress constructs a new APIClient for the server at pachAddr.
func NewFromAddress(pachAddr string) (*APIClient, error) {
	clientConn, err := grpc.Dial(pachAddr, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}

	return &APIClient{
		pfs.NewAPIClient(clientConn),
		pps.NewAPIClient(clientConn),
		pfs.NewBlockAPIClient(clientConn),
	}, nil
}

// NewInCluster constructs a new APIClient using env vars that Kubernetes creates.
// This should be used to access Pachyderm from within a Kubernetes cluster
// with Pachyderm running on it.
func NewInCluster() (*APIClient, error) {
	pachAddr := os.Getenv("PACHD_PORT_650_TCP_ADDR")

	if pachAddr == "" {
		return nil, fmt.Errorf("PACHD_PORT_650_TCP_ADDR not set")
	}

	return NewFromAddress(fmt.Sprintf("%v:650", pachAddr))
}

// DeleteAll deletes everything in the cluster.
// Use with caution, there is no undo.
func (c APIClient) DeleteAll() error {
	if _, err := c.PpsAPIClient.DeleteAll(
		context.Background(),
		google_protobuf.EmptyInstance,
	); err != nil {
		return sanitizeErr(err)
	}
	if _, err := c.PfsAPIClient.DeleteAll(
		context.Background(),
		google_protobuf.EmptyInstance,
	); err != nil {
		return sanitizeErr(err)
	}
	return nil
}

func sanitizeErr(err error) error {
	if err == nil {
		return nil
	}

	return errors.New(grpc.ErrorDesc(err))
}
<commit_msg>Add gRPC connection as property of the ApiClient struct so it can be closed at a later time<commit_after>package client

import (
	"errors"
	"fmt"
	"os"

	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pps"

	google_protobuf "go.pedge.io/pb/go/google/protobuf"
)

// PfsAPIClient is an alias for pfs.APIClient.
type PfsAPIClient pfs.APIClient

// PpsAPIClient is an alias for pps.APIClient.
type PpsAPIClient pps.APIClient

// BlockAPIClient is an alias for pfs.BlockAPIClient.
type BlockAPIClient pfs.BlockAPIClient

// An APIClient is a wrapper around pfs, pps and block APIClients.
type APIClient struct {
	PfsAPIClient
	PpsAPIClient
	BlockAPIClient
	GRPCConn *grpc.ClientConn
}

// NewFromAddress constructs a new APIClient for the server at pachAddr.
func NewFromAddress(pachAddr string) (*APIClient, error) {
	clientConn, err := grpc.Dial(pachAddr, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}

	return &APIClient{
		pfs.NewAPIClient(clientConn),
		pps.NewAPIClient(clientConn),
		pfs.NewBlockAPIClient(clientConn),
		clientConn,
	}, nil
}

// NewInCluster constructs a new APIClient using env vars that Kubernetes creates.
// This should be used to access Pachyderm from within a Kubernetes cluster
// with Pachyderm running on it.
func NewInCluster() (*APIClient, error) {
	pachAddr := os.Getenv("PACHD_PORT_650_TCP_ADDR")

	if pachAddr == "" {
		return nil, fmt.Errorf("PACHD_PORT_650_TCP_ADDR not set")
	}

	return NewFromAddress(fmt.Sprintf("%v:650", pachAddr))
}

// Close the connection to gRPC
func (c *APIClient) Close() error {
	return c.GRPCConn.Close()
}

// DeleteAll deletes everything in the cluster.
// Use with caution, there is no undo.
func (c APIClient) DeleteAll() error {
	if _, err := c.PpsAPIClient.DeleteAll(
		context.Background(),
		google_protobuf.EmptyInstance,
	); err != nil {
		return sanitizeErr(err)
	}
	if _, err := c.PfsAPIClient.DeleteAll(
		context.Background(),
		google_protobuf.EmptyInstance,
	); err != nil {
		return sanitizeErr(err)
	}
	return nil
}

func sanitizeErr(err error) error {
	if err == nil {
		return nil
	}

	return errors.New(grpc.ErrorDesc(err))
}
<|endoftext|>