<commit_before>// Go support for Protocol Buffers - Google's data interchange format
//
// Copyright 2015 The Go Authors.  All rights reserved.
// https://github.com/golang/protobuf
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// Package grpc outputs gRPC service descriptions in Go code.
// It runs as a plugin for the Go protocol buffer compiler plugin.
// It is linked in to protoc-gen-go.
package grpc

import (
	"fmt"
	"path"
	"strconv"
	"strings"

	pb "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/johanbrandhorst/protobuf/protoc-gen-gopherjs/generator"
)

// generatedCodeVersion indicates a version of the generated code.
// It is incremented whenever an incompatibility between the generated code and
// the grpcweb package is introduced; the generated code references
// a constant, grpcweb.GrpcWebPackageIsVersionN (where N is generatedCodeVersion).
const generatedCodeVersion = 2

// Paths for packages used by code generated in this file,
// relative to the import_prefix of the generator.Generator.
const (
	contextPkgPath = "context"
	grpcPkgPath    = "github.com/johanbrandhorst/protobuf/grpcweb"
)

func init() {
	generator.RegisterPlugin(new(grpc))
}

// grpc is an implementation of the Go protocol buffer compiler's
// plugin architecture.  It generates bindings for gRPC support.
type grpc struct {
	gen *generator.Generator
}

// Name returns the name of this plugin, "grpc".
func (g *grpc) Name() string {
	return "grpc"
}

// The names for packages imported in the generated code.
// They may vary from the final path component of the import path
// if the name is used by other packages.
var (
	contextPkg string
	grpcPkg    string
)

// Init initializes the plugin.
func (g *grpc) Init(gen *generator.Generator) {
	g.gen = gen
	contextPkg = generator.RegisterUniquePackageName("context", nil)
	grpcPkg = generator.RegisterUniquePackageName("grpcweb", nil)
}

// Given a type name defined in a .proto, return its object.
// Also record that we're using it, to guarantee the associated import.
func (g *grpc) objectNamed(name string) generator.Object {
	g.gen.RecordTypeUse(name)
	return g.gen.ObjectNamed(name)
}

// Given a type name defined in a .proto, return its name as we will print it.
func (g *grpc) typeName(str string) string {
	return g.gen.TypeName(g.objectNamed(str))
}

// P forwards to g.gen.P.
func (g *grpc) P(args ...interface{}) { g.gen.P(args...) }

// In forwards to g.gen.In.
func (g *grpc) In() { g.gen.In() }

// Out forwards to g.gen.Out.
func (g *grpc) Out() { g.gen.Out() }

// Generate generates code for the services in the given file.
func (g *grpc) Generate(file *generator.FileDescriptor) {
	if len(file.FileDescriptorProto.Service) == 0 {
		return
	}

	g.P("// Reference imports to suppress errors if they are not otherwise used.")
	g.P("var _ ", contextPkg, ".Context")
	g.P("var _ ", grpcPkg, ".Client")
	g.P()

	// Assert version compatibility.
	g.P("// This is a compile-time assertion to ensure that this generated file")
	g.P("// is compatible with the grpcweb package it is being compiled against.")
	g.P("const _ = ", grpcPkg, ".GrpcWebPackageIsVersion", generatedCodeVersion)
	g.P()

	for i, service := range file.FileDescriptorProto.Service {
		g.generateService(file, service, i)
	}
}

// GenerateImports generates the import declaration for this file.
func (g *grpc) GenerateImports(file *generator.FileDescriptor) {
	if len(file.FileDescriptorProto.Service) == 0 {
		return
	}
	g.P("import (")
	g.In()
	g.P(contextPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, contextPkgPath)))
	g.P()
	g.P(grpcPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, grpcPkgPath)))
	g.Out()
	g.P(")")
	g.P()
}

// reservedClientName records whether a client name is reserved on the client side.
var reservedClientName = map[string]bool{
// TODO: do we need any in gRPC?
}

func unexport(s string) string { return strings.ToLower(s[:1]) + s[1:] }

// generateService generates all the code for the named service.
func (g *grpc) generateService(file *generator.FileDescriptor, service *pb.ServiceDescriptorProto, index int) {
	path := fmt.Sprintf("6,%d", index) // 6 means service.

	origServName := service.GetName()
	fullServName := origServName
	if pkg := file.GetPackage(); pkg != "" {
		fullServName = pkg + "." + fullServName
	}
	servName := generator.CamelCase(origServName)

	g.P()
	g.P("// Client API for ", servName, " service")
	g.P()

	// Client interface.
	g.gen.PrintComments(path)
	g.P("type ", servName, "Client interface {")
	g.In()
	for i, method := range service.Method {
		g.gen.PrintComments(fmt.Sprintf("%s,2,%d", path, i)) // 2 means method in a service.
		g.P(g.generateClientSignature(servName, method))
	}
	g.Out()
	g.P("}")
	g.P()

	// Client structure.
	g.P("type ", unexport(servName), "Client struct {")
	g.In()
	g.P("client *", grpcPkg, ".Client")
	g.Out()
	g.P("}")
	g.P()

	// NewClient factory.
	g.P("// New", servName, "Client creates a new gRPC-Web client.")
	g.P("func New", servName, "Client (hostname string, opts ...grpcweb.DialOption) ", servName, "Client {")
	g.In()
	g.P("return &", unexport(servName), "Client{")
	g.In()
	g.P("client: ", grpcPkg, `.NewClient(hostname, "`, fullServName, `", opts...),`)
	g.Out()
	g.P("}")
	g.Out()
	g.P("}")
	g.P()

	serviceDescVar := "_" + servName + "_serviceDesc"
	// Client method implementations.
	for _, method := range service.Method {
		g.generateClientMethod(servName, fullServName, serviceDescVar, method)
	}
}

// generateClientSignature returns the client-side signature for a method.
func (g *grpc) generateClientSignature(servName string, method *pb.MethodDescriptorProto) string {
	origMethName := method.GetName()
	methName := generator.CamelCase(origMethName)
	if reservedClientName[methName] {
		methName += "_"
	}
	reqArg := ", in *" + g.typeName(method.GetInputType())
	if method.GetClientStreaming() {
		reqArg = ""
	}
	respName := "*" + g.typeName(method.GetOutputType())
	if method.GetServerStreaming() || method.GetClientStreaming() {
		respName = servName + "_" + generator.CamelCase(origMethName) + "Client"
	}
	return fmt.Sprintf("%s(ctx %s.Context%s, opts ...%s.CallOption) (%s, error)", methName, contextPkg, reqArg, grpcPkg, respName)
}

func (g *grpc) generateClientMethod(servName, fullServName, serviceDescVar string, method *pb.MethodDescriptorProto) {
	methName := generator.CamelCase(method.GetName())
	outType := g.typeName(method.GetOutputType())
	inType := g.typeName(method.GetInputType())
	streamType := unexport(servName) + methName + "Client"

	g.P("func (c *", unexport(servName), "Client) ", g.generateClientSignature(servName, method), "{")
	g.In()
	switch {
	case !method.GetServerStreaming() && !method.GetClientStreaming():
		// Unary
		g.P(`resp, err := c.client.RPCCall(ctx, "`, method.GetName(), `", in.Marshal(), opts...)`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
		return
	case method.GetServerStreaming() && !method.GetClientStreaming():
		// Server-side stream
		g.P(`srv, err := c.client.NewServerStream(ctx, "`, method.GetName(), `", in.Marshal(), opts...)`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return &", streamType, "{")
		g.In()
		g.P("stream: srv,")
		g.Out()
		g.P("}, nil")
		g.Out()
		g.P("}")
		g.P()
	case method.GetClientStreaming():
		// This case covers both client-side streaming and bidi streaming
		g.P(`srv, err := c.client.NewClientStream(ctx, "`, method.GetName(), `")`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return &", streamType, "{stream: srv}, nil")
		g.Out()
		g.P("}")
		g.P()
	}

	// Stream auxiliary types and methods.
	g.P("type ", servName, "_", methName, "Client interface {")
	g.In()
	if method.GetClientStreaming() {
		g.P("Send(*", inType, ") error")
	}
	if method.GetServerStreaming() {
		g.P("Recv() (*", outType, ", error)")
	}
	if method.GetClientStreaming() && !method.GetServerStreaming() {
		g.P("CloseAndRecv() (*", outType, ", error)")
	} else if method.GetServerStreaming() && method.GetClientStreaming() {
		g.P("CloseSend() error")
	}
	g.P("Context() context.Context")
	g.Out()
	g.P("}")
	g.P()

	g.P("type ", streamType, " struct {")
	g.In()
	if method.GetServerStreaming() && !method.GetClientStreaming() {
		// Server-side streaming only
		g.P("stream ", grpcPkg, ".ServerStream")
	} else {
		// Client side and bidi streams
		g.P("stream ", grpcPkg, ".ClientStream")
	}
	g.Out()
	g.P("}")
	g.P()

	if method.GetClientStreaming() {
		g.P("func (x *", streamType, ") Send(req *", inType, ") error {")
		g.In()
		g.P("return x.stream.SendMsg(req.Marshal())")
		g.Out()
		g.P("}")
		g.P()
	}
	if method.GetServerStreaming() {
		g.P("func (x *", streamType, ") Recv() (*", outType, ", error) {")
		g.In()
		g.P("resp, err := x.stream.RecvMsg()")
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
	}
	if method.GetClientStreaming() && !method.GetServerStreaming() {
		g.P("func (x *", streamType, ") CloseAndRecv() (*", outType, ", error) {")
		g.In()
		g.P("resp, err := x.stream.CloseAndRecv()")
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
	} else if method.GetServerStreaming() && method.GetClientStreaming() {
		g.P("func (x *", streamType, ") CloseSend() error {")
		g.In()
		g.P("return x.stream.CloseSend()")
		g.Out()
		g.P("}")
		g.P()
	}
	g.P("func (x *", streamType, ") Context() context.Context {")
	g.In()
	g.P("return x.Context()")
	g.Out()
	g.P("}")
	g.P()
}
<commit_msg>Correct Context function<commit_after>// Go support for Protocol Buffers - Google's data interchange format
//
// Copyright 2015 The Go Authors.  All rights reserved.
// https://github.com/golang/protobuf
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// Package grpc outputs gRPC service descriptions in Go code.
// It runs as a plugin for the Go protocol buffer compiler plugin.
// It is linked in to protoc-gen-go.
package grpc

import (
	"fmt"
	"path"
	"strconv"
	"strings"

	pb "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/johanbrandhorst/protobuf/protoc-gen-gopherjs/generator"
)

// generatedCodeVersion indicates a version of the generated code.
// It is incremented whenever an incompatibility between the generated code and
// the grpcweb package is introduced; the generated code references
// a constant, grpcweb.GrpcWebPackageIsVersionN (where N is generatedCodeVersion).
const generatedCodeVersion = 2

// Paths for packages used by code generated in this file,
// relative to the import_prefix of the generator.Generator.
const (
	contextPkgPath = "context"
	grpcPkgPath    = "github.com/johanbrandhorst/protobuf/grpcweb"
)

func init() {
	generator.RegisterPlugin(new(grpc))
}

// grpc is an implementation of the Go protocol buffer compiler's
// plugin architecture.  It generates bindings for gRPC support.
type grpc struct {
	gen *generator.Generator
}

// Name returns the name of this plugin, "grpc".
func (g *grpc) Name() string {
	return "grpc"
}

// The names for packages imported in the generated code.
// They may vary from the final path component of the import path
// if the name is used by other packages.
var (
	contextPkg string
	grpcPkg    string
)

// Init initializes the plugin.
func (g *grpc) Init(gen *generator.Generator) {
	g.gen = gen
	contextPkg = generator.RegisterUniquePackageName("context", nil)
	grpcPkg = generator.RegisterUniquePackageName("grpcweb", nil)
}

// Given a type name defined in a .proto, return its object.
// Also record that we're using it, to guarantee the associated import.
func (g *grpc) objectNamed(name string) generator.Object {
	g.gen.RecordTypeUse(name)
	return g.gen.ObjectNamed(name)
}

// Given a type name defined in a .proto, return its name as we will print it.
func (g *grpc) typeName(str string) string {
	return g.gen.TypeName(g.objectNamed(str))
}

// P forwards to g.gen.P.
func (g *grpc) P(args ...interface{}) { g.gen.P(args...) }

// In forwards to g.gen.In.
func (g *grpc) In() { g.gen.In() }

// Out forwards to g.gen.Out.
func (g *grpc) Out() { g.gen.Out() }

// Generate generates code for the services in the given file.
func (g *grpc) Generate(file *generator.FileDescriptor) {
	if len(file.FileDescriptorProto.Service) == 0 {
		return
	}

	g.P("// Reference imports to suppress errors if they are not otherwise used.")
	g.P("var _ ", contextPkg, ".Context")
	g.P("var _ ", grpcPkg, ".Client")
	g.P()

	// Assert version compatibility.
	g.P("// This is a compile-time assertion to ensure that this generated file")
	g.P("// is compatible with the grpcweb package it is being compiled against.")
	g.P("const _ = ", grpcPkg, ".GrpcWebPackageIsVersion", generatedCodeVersion)
	g.P()

	for i, service := range file.FileDescriptorProto.Service {
		g.generateService(file, service, i)
	}
}

// GenerateImports generates the import declaration for this file.
func (g *grpc) GenerateImports(file *generator.FileDescriptor) {
	if len(file.FileDescriptorProto.Service) == 0 {
		return
	}
	g.P("import (")
	g.In()
	g.P(contextPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, contextPkgPath)))
	g.P()
	g.P(grpcPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, grpcPkgPath)))
	g.Out()
	g.P(")")
	g.P()
}

// reservedClientName records whether a client name is reserved on the client side.
var reservedClientName = map[string]bool{
// TODO: do we need any in gRPC?
}

func unexport(s string) string { return strings.ToLower(s[:1]) + s[1:] }

// generateService generates all the code for the named service.
func (g *grpc) generateService(file *generator.FileDescriptor, service *pb.ServiceDescriptorProto, index int) {
	path := fmt.Sprintf("6,%d", index) // 6 means service.

	origServName := service.GetName()
	fullServName := origServName
	if pkg := file.GetPackage(); pkg != "" {
		fullServName = pkg + "." + fullServName
	}
	servName := generator.CamelCase(origServName)

	g.P()
	g.P("// Client API for ", servName, " service")
	g.P()

	// Client interface.
	g.gen.PrintComments(path)
	g.P("type ", servName, "Client interface {")
	g.In()
	for i, method := range service.Method {
		g.gen.PrintComments(fmt.Sprintf("%s,2,%d", path, i)) // 2 means method in a service.
		g.P(g.generateClientSignature(servName, method))
	}
	g.Out()
	g.P("}")
	g.P()

	// Client structure.
	g.P("type ", unexport(servName), "Client struct {")
	g.In()
	g.P("client *", grpcPkg, ".Client")
	g.Out()
	g.P("}")
	g.P()

	// NewClient factory.
	g.P("// New", servName, "Client creates a new gRPC-Web client.")
	g.P("func New", servName, "Client (hostname string, opts ...grpcweb.DialOption) ", servName, "Client {")
	g.In()
	g.P("return &", unexport(servName), "Client{")
	g.In()
	g.P("client: ", grpcPkg, `.NewClient(hostname, "`, fullServName, `", opts...),`)
	g.Out()
	g.P("}")
	g.Out()
	g.P("}")
	g.P()

	serviceDescVar := "_" + servName + "_serviceDesc"
	// Client method implementations.
	for _, method := range service.Method {
		g.generateClientMethod(servName, fullServName, serviceDescVar, method)
	}
}

// generateClientSignature returns the client-side signature for a method.
func (g *grpc) generateClientSignature(servName string, method *pb.MethodDescriptorProto) string {
	origMethName := method.GetName()
	methName := generator.CamelCase(origMethName)
	if reservedClientName[methName] {
		methName += "_"
	}
	reqArg := ", in *" + g.typeName(method.GetInputType())
	if method.GetClientStreaming() {
		reqArg = ""
	}
	respName := "*" + g.typeName(method.GetOutputType())
	if method.GetServerStreaming() || method.GetClientStreaming() {
		respName = servName + "_" + generator.CamelCase(origMethName) + "Client"
	}
	return fmt.Sprintf("%s(ctx %s.Context%s, opts ...%s.CallOption) (%s, error)", methName, contextPkg, reqArg, grpcPkg, respName)
}

func (g *grpc) generateClientMethod(servName, fullServName, serviceDescVar string, method *pb.MethodDescriptorProto) {
	methName := generator.CamelCase(method.GetName())
	outType := g.typeName(method.GetOutputType())
	inType := g.typeName(method.GetInputType())
	streamType := unexport(servName) + methName + "Client"

	g.P("func (c *", unexport(servName), "Client) ", g.generateClientSignature(servName, method), "{")
	g.In()
	switch {
	case !method.GetServerStreaming() && !method.GetClientStreaming():
		// Unary
		g.P(`resp, err := c.client.RPCCall(ctx, "`, method.GetName(), `", in.Marshal(), opts...)`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
		return
	case method.GetServerStreaming() && !method.GetClientStreaming():
		// Server-side stream
		g.P(`srv, err := c.client.NewServerStream(ctx, "`, method.GetName(), `", in.Marshal(), opts...)`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return &", streamType, "{")
		g.In()
		g.P("stream: srv,")
		g.Out()
		g.P("}, nil")
		g.Out()
		g.P("}")
		g.P()
	case method.GetClientStreaming():
		// This case covers both client-side streaming and bidi streaming
		g.P(`srv, err := c.client.NewClientStream(ctx, "`, method.GetName(), `")`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return &", streamType, "{stream: srv}, nil")
		g.Out()
		g.P("}")
		g.P()
	}

	// Stream auxiliary types and methods.
	g.P("type ", servName, "_", methName, "Client interface {")
	g.In()
	if method.GetClientStreaming() {
		g.P("Send(*", inType, ") error")
	}
	if method.GetServerStreaming() {
		g.P("Recv() (*", outType, ", error)")
	}
	if method.GetClientStreaming() && !method.GetServerStreaming() {
		g.P("CloseAndRecv() (*", outType, ", error)")
	} else if method.GetServerStreaming() && method.GetClientStreaming() {
		g.P("CloseSend() error")
	}
	g.P("Context() context.Context")
	g.Out()
	g.P("}")
	g.P()

	g.P("type ", streamType, " struct {")
	g.In()
	if method.GetServerStreaming() && !method.GetClientStreaming() {
		// Server-side streaming only
		g.P("stream ", grpcPkg, ".ServerStream")
	} else {
		// Client side and bidi streams
		g.P("stream ", grpcPkg, ".ClientStream")
	}
	g.Out()
	g.P("}")
	g.P()

	if method.GetClientStreaming() {
		g.P("func (x *", streamType, ") Send(req *", inType, ") error {")
		g.In()
		g.P("return x.stream.SendMsg(req.Marshal())")
		g.Out()
		g.P("}")
		g.P()
	}
	if method.GetServerStreaming() {
		g.P("func (x *", streamType, ") Recv() (*", outType, ", error) {")
		g.In()
		g.P("resp, err := x.stream.RecvMsg()")
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
	}
	if method.GetClientStreaming() && !method.GetServerStreaming() {
		g.P("func (x *", streamType, ") CloseAndRecv() (*", outType, ", error) {")
		g.In()
		g.P("resp, err := x.stream.CloseAndRecv()")
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
	} else if method.GetServerStreaming() && method.GetClientStreaming() {
		g.P("func (x *", streamType, ") CloseSend() error {")
		g.In()
		g.P("return x.stream.CloseSend()")
		g.Out()
		g.P("}")
		g.P()
	}
	g.P("func (x *", streamType, ") Context() context.Context {")
	g.In()
	g.P("return x.stream.Context()")
	g.Out()
	g.P("}")
	g.P()
}
<|endoftext|>


<commit_before>// This package implements a provisioner for Packer that executes
// shell scripts within the remote machine.
package shell

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// If true, the script contains binary and line endings will not be
	// converted from Windows to Unix-style.
	Binary bool

	// An inline script to execute. Multiple strings are all executed
	// in the context of a single shell.
	Inline []string

	// The shebang value used when running inline scripts.
	InlineShebang string `mapstructure:"inline_shebang"`

	// The local path of the shell script to upload and execute.
	Script string

	// An array of multiple scripts to run.
	Scripts []string

	// An array of environment variables that will be injected before
	// your command(s) are executed.
	Vars []string `mapstructure:"environment_vars"`

	// The remote path where the local shell script will be uploaded to.
	// This should be set to a writable file that is in a pre-existing directory.
	RemotePath string `mapstructure:"remote_path"`

	// The command used to execute the script. The '{{ .Path }}' variable
	// should be used to specify where the script goes, {{ .Vars }}
	// can be used to inject the environment_vars into the environment.
	ExecuteCommand string `mapstructure:"execute_command"`

	// The timeout for retrying to start the process. Until this timeout
	// is reached, if the provisioner can't start a process, it retries.
	// This can be set high to allow for reboots.
	RawStartRetryTimeout string `mapstructure:"start_retry_timeout"`

	// Whether to clean scripts up
	NoClean bool

	startRetryTimeout time.Duration
	ctx               interpolate.Context
}

type Provisioner struct {
	config Config
}

type ExecuteCommandTemplate struct {
	Vars string
	Path string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "chmod +x {{.Path}}; {{.Vars}} {{.Path}}"
	}

	if p.config.Inline != nil && len(p.config.Inline) == 0 {
		p.config.Inline = nil
	}

	if p.config.InlineShebang == "" {
		p.config.InlineShebang = "/bin/sh -e"
	}

	if p.config.RawStartRetryTimeout == "" {
		p.config.RawStartRetryTimeout = "5m"
	}

	if p.config.RemotePath == "" {
		p.config.RemotePath = fmt.Sprintf(
			"/tmp/script_%d.sh", rand.Intn(9999))
	}

	if p.config.Scripts == nil {
		p.config.Scripts = make([]string, 0)
	}

	if p.config.Vars == nil {
		p.config.Vars = make([]string, 0)
	}

	var errs *packer.MultiError
	if p.config.Script != "" && len(p.config.Scripts) > 0 {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only one of script or scripts can be specified."))
	}

	if p.config.Script != "" {
		p.config.Scripts = []string{p.config.Script}
	}

	if len(p.config.Scripts) == 0 && p.config.Inline == nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Either a script file or inline script must be specified."))
	} else if len(p.config.Scripts) > 0 && p.config.Inline != nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only a script file or an inline script can be specified, not both."))
	}

	for _, path := range p.config.Scripts {
		if _, err := os.Stat(path); err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Bad script '%s': %s", path, err))
		}
	}

	// Do a check for bad environment variables, such as '=foo', 'foobar'
	for idx, kv := range p.config.Vars {
		vs := strings.SplitN(kv, "=", 2)
		if len(vs) != 2 || vs[0] == "" {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
		} else {
			// Replace single quotes so they parse
			vs[1] = strings.Replace(vs[1], "'", `'"'"'`, -1)

			// Single quote env var values
			p.config.Vars[idx] = fmt.Sprintf("%s='%s'", vs[0], vs[1])
		}
	}

	if p.config.RawStartRetryTimeout != "" {
		p.config.startRetryTimeout, err = time.ParseDuration(p.config.RawStartRetryTimeout)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Failed parsing start_retry_timeout: %s", err))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	scripts := make([]string, len(p.config.Scripts))
	copy(scripts, p.config.Scripts)

	// If we have an inline script, then turn that into a temporary
	// shell script and use that.
	if p.config.Inline != nil {
		tf, err := ioutil.TempFile("", "packer-shell")
		if err != nil {
			return fmt.Errorf("Error preparing shell script: %s", err)
		}
		defer os.Remove(tf.Name())

		// Set the path to the temporary file
		scripts = append(scripts, tf.Name())

		// Write our contents to it
		writer := bufio.NewWriter(tf)
		writer.WriteString(fmt.Sprintf("#!%s\n", p.config.InlineShebang))
		for _, command := range p.config.Inline {
			if _, err := writer.WriteString(command + "\n"); err != nil {
				return fmt.Errorf("Error preparing shell script: %s", err)
			}
		}

		if err := writer.Flush(); err != nil {
			return fmt.Errorf("Error preparing shell script: %s", err)
		}

		tf.Close()
	}

	// Build our variables up by adding in the build name and builder type
	envVars := make([]string, len(p.config.Vars)+2)
	envVars[0] = fmt.Sprintf("PACKER_BUILD_NAME='%s'", p.config.PackerBuildName)
	envVars[1] = fmt.Sprintf("PACKER_BUILDER_TYPE='%s'", p.config.PackerBuilderType)
	copy(envVars[2:], p.config.Vars)

	for _, path := range scripts {
		ui.Say(fmt.Sprintf("Provisioning with shell script: %s", path))

		log.Printf("Opening %s for reading", path)
		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Error opening shell script: %s", err)
		}
		defer f.Close()

		// Flatten the environment variables
		flattendVars := strings.Join(envVars, " ")

		// Compile the command
		p.config.ctx.Data = &ExecuteCommandTemplate{
			Vars: flattendVars,
			Path: p.config.RemotePath,
		}
		command, err := interpolate.Render(p.config.ExecuteCommand, &p.config.ctx)
		if err != nil {
			return fmt.Errorf("Error processing command: %s", err)
		}

		// Upload the file and run the command. Do this in the context of
		// a single retryable function so that we don't end up with
		// the case that the upload succeeded, a restart is initiated,
		// and then the command is executed but the file doesn't exist
		// any longer.
		var cmd *packer.RemoteCmd
		err = p.retryable(func() error {
			if _, err := f.Seek(0, 0); err != nil {
				return err
			}

			var r io.Reader = f
			if !p.config.Binary {
				r = &UnixReader{Reader: r}
			}

			if err := comm.Upload(p.config.RemotePath, r, nil); err != nil {
				return fmt.Errorf("Error uploading script: %s", err)
			}

			cmd = &packer.RemoteCmd{
				Command: fmt.Sprintf("chmod 0755 %s", p.config.RemotePath),
			}
			if err := comm.Start(cmd); err != nil {
				return fmt.Errorf(
					"Error chmodding script file to 0755 in remote "+
						"machine: %s", err)
			}
			cmd.Wait()

			cmd = &packer.RemoteCmd{Command: command}
			return cmd.StartWithUi(comm, ui)
		})
		if err != nil {
			return err
		}

		if cmd.ExitStatus != 0 {
			return fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus)
		}

		if !p.config.NoClean {

			// Delete the temporary file we created. We retry this a few times
			// since if the above rebooted we have to wait until the reboot
			// completes.
			err = p.retryable(func() error {
				cmd = &packer.RemoteCmd{
					Command: fmt.Sprintf("rm -f %s", p.config.RemotePath),
				}
				if err := comm.Start(cmd); err != nil {
					return fmt.Errorf(
						"Error removing temporary script at %s: %s",
						p.config.RemotePath, err)
				}
				cmd.Wait()
				return nil
			})
			if err != nil {
				return err
			}

			if cmd.ExitStatus != 0 {
				return fmt.Errorf(
					"Error removing temporary script at %s!",
					p.config.RemotePath)
			}
		}
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

// retryable will retry the given function over and over until a
// non-error is returned.
func (p *Provisioner) retryable(f func() error) error {
	startTimeout := time.After(p.config.startRetryTimeout)
	for {
		var err error
		if err = f(); err == nil {
			return nil
		}

		// Create an error and log it
		err = fmt.Errorf("Retryable error: %s", err)
		log.Printf(err.Error())

		// Check if we timed out, otherwise we retry. It is safe to
		// retry since the only error case above is if the command
		// failed to START.
		select {
		case <-startTimeout:
			return err
		default:
			time.Sleep(2 * time.Second)
		}
	}
}
<commit_msg>Rename noclean shell option to skip_clean, per pull feedback<commit_after>// This package implements a provisioner for Packer that executes
// shell scripts within the remote machine.
package shell

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// If true, the script contains binary and line endings will not be
	// converted from Windows to Unix-style.
	Binary bool

	// An inline script to execute. Multiple strings are all executed
	// in the context of a single shell.
	Inline []string

	// The shebang value used when running inline scripts.
	InlineShebang string `mapstructure:"inline_shebang"`

	// The local path of the shell script to upload and execute.
	Script string

	// An array of multiple scripts to run.
	Scripts []string

	// An array of environment variables that will be injected before
	// your command(s) are executed.
	Vars []string `mapstructure:"environment_vars"`

	// The remote path where the local shell script will be uploaded to.
	// This should be set to a writable file that is in a pre-existing directory.
	RemotePath string `mapstructure:"remote_path"`

	// The command used to execute the script. The '{{ .Path }}' variable
	// should be used to specify where the script goes, {{ .Vars }}
	// can be used to inject the environment_vars into the environment.
	ExecuteCommand string `mapstructure:"execute_command"`

	// The timeout for retrying to start the process. Until this timeout
	// is reached, if the provisioner can't start a process, it retries.
	// This can be set high to allow for reboots.
	RawStartRetryTimeout string `mapstructure:"start_retry_timeout"`

	// Whether to clean scripts up
	SkipClean bool `mapstructure:"skip_clean"`

	startRetryTimeout time.Duration
	ctx               interpolate.Context
}

type Provisioner struct {
	config Config
}

type ExecuteCommandTemplate struct {
	Vars string
	Path string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "chmod +x {{.Path}}; {{.Vars}} {{.Path}}"
	}

	if p.config.Inline != nil && len(p.config.Inline) == 0 {
		p.config.Inline = nil
	}

	if p.config.InlineShebang == "" {
		p.config.InlineShebang = "/bin/sh -e"
	}

	if p.config.RawStartRetryTimeout == "" {
		p.config.RawStartRetryTimeout = "5m"
	}

	if p.config.RemotePath == "" {
		p.config.RemotePath = fmt.Sprintf(
			"/tmp/script_%d.sh", rand.Intn(9999))
	}

	if p.config.Scripts == nil {
		p.config.Scripts = make([]string, 0)
	}

	if p.config.Vars == nil {
		p.config.Vars = make([]string, 0)
	}

	var errs *packer.MultiError
	if p.config.Script != "" && len(p.config.Scripts) > 0 {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only one of script or scripts can be specified."))
	}

	if p.config.Script != "" {
		p.config.Scripts = []string{p.config.Script}
	}

	if len(p.config.Scripts) == 0 && p.config.Inline == nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Either a script file or inline script must be specified."))
	} else if len(p.config.Scripts) > 0 && p.config.Inline != nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only a script file or an inline script can be specified, not both."))
	}

	for _, path := range p.config.Scripts {
		if _, err := os.Stat(path); err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Bad script '%s': %s", path, err))
		}
	}

	// Do a check for bad environment variables, such as '=foo', 'foobar'
	for idx, kv := range p.config.Vars {
		vs := strings.SplitN(kv, "=", 2)
		if len(vs) != 2 || vs[0] == "" {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
		} else {
			// Replace single quotes so they parse
			vs[1] = strings.Replace(vs[1], "'", `'"'"'`, -1)

			// Single quote env var values
			p.config.Vars[idx] = fmt.Sprintf("%s='%s'", vs[0], vs[1])
		}
	}

	if p.config.RawStartRetryTimeout != "" {
		p.config.startRetryTimeout, err = time.ParseDuration(p.config.RawStartRetryTimeout)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Failed parsing start_retry_timeout: %s", err))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	scripts := make([]string, len(p.config.Scripts))
	copy(scripts, p.config.Scripts)

	// If we have an inline script, then turn that into a temporary
	// shell script and use that.
	if p.config.Inline != nil {
		tf, err := ioutil.TempFile("", "packer-shell")
		if err != nil {
			return fmt.Errorf("Error preparing shell script: %s", err)
		}
		defer os.Remove(tf.Name())

		// Set the path to the temporary file
		scripts = append(scripts, tf.Name())

		// Write our contents to it
		writer := bufio.NewWriter(tf)
		writer.WriteString(fmt.Sprintf("#!%s\n", p.config.InlineShebang))
		for _, command := range p.config.Inline {
			if _, err := writer.WriteString(command + "\n"); err != nil {
				return fmt.Errorf("Error preparing shell script: %s", err)
			}
		}

		if err := writer.Flush(); err != nil {
			return fmt.Errorf("Error preparing shell script: %s", err)
		}

		tf.Close()
	}

	// Build our variables up by adding in the build name and builder type
	envVars := make([]string, len(p.config.Vars)+2)
	envVars[0] = fmt.Sprintf("PACKER_BUILD_NAME='%s'", p.config.PackerBuildName)
	envVars[1] = fmt.Sprintf("PACKER_BUILDER_TYPE='%s'", p.config.PackerBuilderType)
	copy(envVars[2:], p.config.Vars)

	for _, path := range scripts {
		ui.Say(fmt.Sprintf("Provisioning with shell script: %s", path))

		log.Printf("Opening %s for reading", path)
		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Error opening shell script: %s", err)
		}
		defer f.Close()

		// Flatten the environment variables
		flattendVars := strings.Join(envVars, " ")

		// Compile the command
		p.config.ctx.Data = &ExecuteCommandTemplate{
			Vars: flattendVars,
			Path: p.config.RemotePath,
		}
		command, err := interpolate.Render(p.config.ExecuteCommand, &p.config.ctx)
		if err != nil {
			return fmt.Errorf("Error processing command: %s", err)
		}

		// Upload the file and run the command. Do this in the context of
		// a single retryable function so that we don't end up with
		// the case that the upload succeeded, a restart is initiated,
		// and then the command is executed but the file doesn't exist
		// any longer.
		var cmd *packer.RemoteCmd
		err = p.retryable(func() error {
			if _, err := f.Seek(0, 0); err != nil {
				return err
			}

			var r io.Reader = f
			if !p.config.Binary {
				r = &UnixReader{Reader: r}
			}

			if err := comm.Upload(p.config.RemotePath, r, nil); err != nil {
				return fmt.Errorf("Error uploading script: %s", err)
			}

			cmd = &packer.RemoteCmd{
				Command: fmt.Sprintf("chmod 0755 %s", p.config.RemotePath),
			}
			if err := comm.Start(cmd); err != nil {
				return fmt.Errorf(
					"Error chmodding script file to 0755 in remote "+
						"machine: %s", err)
			}
			cmd.Wait()

			cmd = &packer.RemoteCmd{Command: command}
			return cmd.StartWithUi(comm, ui)
		})
		if err != nil {
			return err
		}

		if cmd.ExitStatus != 0 {
			return fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus)
		}

		if !p.config.SkipClean {

			// Delete the temporary file we created. We retry this a few times
			// since if the above rebooted we have to wait until the reboot
			// completes.
			err = p.retryable(func() error {
				cmd = &packer.RemoteCmd{
					Command: fmt.Sprintf("rm -f %s", p.config.RemotePath),
				}
				if err := comm.Start(cmd); err != nil {
					return fmt.Errorf(
						"Error removing temporary script at %s: %s",
						p.config.RemotePath, err)
				}
				cmd.Wait()
				return nil
			})
			if err != nil {
				return err
			}

			if cmd.ExitStatus != 0 {
				return fmt.Errorf(
					"Error removing temporary script at %s!",
					p.config.RemotePath)
			}
		}
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

// retryable will retry the given function over and over until a
// non-error is returned.
func (p *Provisioner) retryable(f func() error) error {
	startTimeout := time.After(p.config.startRetryTimeout)
	for {
		var err error
		if err = f(); err == nil {
			return nil
		}

		// Create an error and log it
		err = fmt.Errorf("Retryable error: %s", err)
		log.Printf(err.Error())

		// Check if we timed out, otherwise we retry. It is safe to
		// retry since the only error case above is if the command
		// failed to START.
		select {
		case <-startTimeout:
			return err
		default:
			time.Sleep(2 * time.Second)
		}
	}
}
<|endoftext|>


<commit_before>package udp_tracker

import (
	"bitbucket.org/anacrolix/go.torrent/tracker"
	"bytes"
	"encoding/binary"
	"io"
	"math/rand"
	"net"
	"net/url"
	"time"
)

type Action int32

const (
	Connect Action = iota
	Announce
	Scrape
	Error
)

type ConnectionRequest struct {
	ConnectionId int64
	Action       int32
	TransctionId int32
}

type ConnectionResponse struct {
	ConnectionId int64
}

type ResponseHeader struct {
	Action        Action
	TransactionId int32
}

type RequestHeader struct {
	ConnectionId  int64
	Action        Action
	TransactionId int32
}

type AnnounceResponseHeader struct {
	Interval int32
	Leechers int32
	Seeders  int32
}

type Peer struct {
	IP   [4]byte
	Port uint16
}

func init() {
	tracker.RegisterClientScheme("udp", newClient)
}

func newClient(url *url.URL) tracker.Client {
	return &client{}
}

func newTransactionId() int32 {
	return int32(rand.Uint32())
}

func timeout(contiguousTimeouts int) (d time.Duration) {
	if contiguousTimeouts > 8 {
		contiguousTimeouts = 8
	}
	d = 15 * time.Second
	for ; contiguousTimeouts > 0; contiguousTimeouts-- {
		d *= 2
	}
	return
}

type client struct {
	contiguousTimeouts   int
	connectionIdReceived time.Time
	connectionId         int64
	socket               net.Conn
}

func (c *client) Announce(req *tracker.AnnounceRequest) (res tracker.AnnounceResponse, err error) {
	err = c.connect()
	if err != nil {
		return
	}
	b, err := c.request(Announce, req)
	if err != nil {
		return
	}
	var (
		h  AnnounceResponseHeader
		ps []Peer
	)
	err = readBody(b, &h, &ps)
	if err != nil {
		return
	}
	res.Interval = h.Interval
	res.Leechers = h.Leechers
	res.Seeders = h.Seeders
	for _, p := range ps {
		res.Peers = append(res.Peers, tracker.Peer{
			IP:   p.IP[:],
			Port: int(p.Port),
		})
	}
	return
}

func (c *client) write(h *RequestHeader, body interface{}) (err error) {
	buf := &bytes.Buffer{}
	err = binary.Write(buf, binary.BigEndian, h)
	if err != nil {
		panic(err)
	}
	err = binary.Write(buf, binary.BigEndian, body)
	if err != nil {
		panic(err)
	}
	n, err := c.socket.Write(buf.Bytes())
	if err != nil {
		return
	}
	if n != buf.Len() {
		panic("write should send all or error")
	}
	return
}

func (c *client) request(action Action, args interface{}) (responseBody []byte, err error) {
	tid := newTransactionId()
	err = c.write(&RequestHeader{
		ConnectionId:  c.connectionId,
		Action:        action,
		TransactionId: tid,
	}, args)
	if err != nil {
		return
	}
	c.socket.SetDeadline(time.Now().Add(timeout(c.contiguousTimeouts)))
	b := make([]byte, 0x10000) // IP limits packet size to 64KB
	for {
		var n int
		n, err = c.socket.Read(b)
		if opE, ok := err.(*net.OpError); ok {
			if opE.Timeout() {
				c.contiguousTimeouts++
				return
			}
		}
		if err != nil {
			return
		}
		buf := bytes.NewBuffer(b[:n])
		var h ResponseHeader
		err = binary.Read(buf, binary.BigEndian, &h)
		switch err {
		case io.ErrUnexpectedEOF:
			continue
		case nil:
		default:
			return
		}
		if h.Action != action {
			continue
		}
		if h.TransactionId != tid {
			continue
		}
		c.contiguousTimeouts = 0
		responseBody = buf.Bytes()
		return
	}
}

func readBody(b []byte, data ...interface{}) (err error) {
	r := bytes.NewReader(b)
	for _, datum := range data {
		err = binary.Read(r, binary.BigEndian, datum)
		if err != nil {
			break
		}
	}
	return
}

func (c *client) connect() (err error) {
	if !c.connectionIdReceived.IsZero() && time.Now().Before(c.connectionIdReceived.Add(time.Minute)) {
		return nil
	}
	c.connectionId = 0x41727101980
	b, err := c.request(Connect, nil)
	if err != nil {
		return
	}
	var res ConnectionResponse
	err = readBody(b, &res)
	if err != nil {
		return
	}
	c.connectionId = res.ConnectionId
	c.connectionIdReceived = time.Now()
	return
}
<commit_msg>Fix parsing of UDP announce response peers<commit_after>package udp_tracker

import (
	"bitbucket.org/anacrolix/go.torrent/tracker"
	"bytes"
	"encoding/binary"
	"io"
	"math/rand"
	"net"
	"net/url"
	"time"
)

type Action int32

const (
	Connect Action = iota
	Announce
	Scrape
	Error
)

type ConnectionRequest struct {
	ConnectionId int64
	Action       int32
	TransctionId int32
}

type ConnectionResponse struct {
	ConnectionId int64
}

type ResponseHeader struct {
	Action        Action
	TransactionId int32
}

type RequestHeader struct {
	ConnectionId  int64
	Action        Action
	TransactionId int32
}

type AnnounceResponseHeader struct {
	Interval int32
	Leechers int32
	Seeders  int32
}

type Peer struct {
	IP   [4]byte
	Port uint16
}

func init() {
	tracker.RegisterClientScheme("udp", newClient)
}

func newClient(url *url.URL) tracker.Client {
	return &client{}
}

func newTransactionId() int32 {
	return int32(rand.Uint32())
}

func timeout(contiguousTimeouts int) (d time.Duration) {
	if contiguousTimeouts > 8 {
		contiguousTimeouts = 8
	}
	d = 15 * time.Second
	for ; contiguousTimeouts > 0; contiguousTimeouts-- {
		d *= 2
	}
	return
}

type client struct {
	contiguousTimeouts   int
	connectionIdReceived time.Time
	connectionId         int64
	socket               net.Conn
}

func (c *client) Announce(req *tracker.AnnounceRequest) (res tracker.AnnounceResponse, err error) {
	err = c.connect()
	if err != nil {
		return
	}
	b, err := c.request(Announce, req)
	if err != nil {
		return
	}
	var h AnnounceResponseHeader
	err = readBody(b, &h)
	if err != nil {
		return
	}
	res.Interval = h.Interval
	res.Leechers = h.Leechers
	res.Seeders = h.Seeders
	for {
		var p Peer
		err = binary.Read(b, binary.BigEndian, &p)
		switch err {
		case nil:
		case io.EOF:
			err = nil
			fallthrough
		default:
			return
		}
		res.Peers = append(res.Peers, tracker.Peer{
			IP:   p.IP[:],
			Port: int(p.Port),
		})
	}
	return
}

func (c *client) write(h *RequestHeader, body interface{}) (err error) {
	buf := &bytes.Buffer{}
	err = binary.Write(buf, binary.BigEndian, h)
	if err != nil {
		panic(err)
	}
	err = binary.Write(buf, binary.BigEndian, body)
	if err != nil {
		panic(err)
	}
	n, err := c.socket.Write(buf.Bytes())
	if err != nil {
		return
	}
	if n != buf.Len() {
		panic("write should send all or error")
	}
	return
}

func (c *client) request(action Action, args interface{}) (responseBody *bytes.Reader, err error) {
	tid := newTransactionId()
	err = c.write(&RequestHeader{
		ConnectionId:  c.connectionId,
		Action:        action,
		TransactionId: tid,
	}, args)
	if err != nil {
		return
	}
	c.socket.SetDeadline(time.Now().Add(timeout(c.contiguousTimeouts)))
	b := make([]byte, 0x10000) // IP limits packet size to 64KB
	for {
		var n int
		n, err = c.socket.Read(b)
		if opE, ok := err.(*net.OpError); ok {
			if opE.Timeout() {
				c.contiguousTimeouts++
				return
			}
		}
		if err != nil {
			return
		}
		buf := bytes.NewBuffer(b[:n])
		var h ResponseHeader
		err = binary.Read(buf, binary.BigEndian, &h)
		switch err {
		case io.ErrUnexpectedEOF:
			continue
		case nil:
		default:
			return
		}
		if h.Action != action {
			continue
		}
		if h.TransactionId != tid {
			continue
		}
		c.contiguousTimeouts = 0
		if h.Action == Error {
			err = errors.New(buf.String())
		}
		responseBody = bytes.NewReader(buf.Bytes())
		return
	}
}

func readBody(r *bytes.Reader, data ...interface{}) (err error) {
	for _, datum := range data {
		err = binary.Read(r, binary.BigEndian, datum)
		if err != nil {
			break
		}
	}
	return
}

func (c *client) connect() (err error) {
	if !c.connectionIdReceived.IsZero() && time.Now().Before(c.connectionIdReceived.Add(time.Minute)) {
		return nil
	}
	c.connectionId = 0x41727101980
	b, err := c.request(Connect, nil)
	if err != nil {
		return
	}
	var res ConnectionResponse
	err = readBody(b, &res)
	if err != nil {
		return
	}
	c.connectionId = res.ConnectionId
	c.connectionIdReceived = time.Now()
	return
}
<|endoftext|>


<commit_before>package transformer

var (
	ValidTransforms = []string{
		"bigint",
		"bool",
		"float",
		"int",
		"ipAsn",
		"ipAsnInteger",
		"ipCity",
		"ipCountry",
		"ipRegion",
		"stringToIntegerMD5",
		"varchar",
		"f@timestamp@unix",
	}
)
<commit_msg>Add f@timestamp@unix-utc<commit_after>package transformer

var (
	ValidTransforms = []string{
		"bigint",
		"bool",
		"float",
		"int",
		"ipAsn",
		"ipAsnInteger",
		"ipCity",
		"ipCountry",
		"ipRegion",
		"stringToIntegerMD5",
		"varchar",
		"f@timestamp@unix",
		"f@timestamp@unix-utc",
	}
)
<|endoftext|>


<commit_before>package abide

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"strings"
	"testing"

	"github.com/beme/abide/internal"
	"github.com/sergi/go-diff/diffmatchpatch"
)

// Assertable represents an object that can be asserted.
type Assertable interface {
	String() string
}

// Assert asserts the value of an object with implements Assertable.
func Assert(t *testing.T, id string, a Assertable) {
	data := a.String()
	createOrUpdateSnapshot(t, id, data)
}

// AssertHTTPResponse asserts the value of an http.Response.
func AssertHTTPResponse(t *testing.T, id string, w *http.Response) {
	body, err := httputil.DumpResponse(w, true)
	if err != nil {
		t.Fatal(err)
	}

	assertHTTP(t, id, body, contentTypeIsJSON(w.Header.Get("Content-Type")))
}

// AssertHTTPRequestOut asserts the value of an http.Request.
// Intended for use when testing outgoing client requests
// See https://golang.org/pkg/net/http/httputil/#DumpRequestOut for more
func AssertHTTPRequestOut(t *testing.T, id string, r *http.Request) {
	body, err := httputil.DumpRequestOut(r, true)
	if err != nil {
		t.Fatal(err)
	}

	assertHTTP(t, id, body, contentTypeIsJSON(r.Header.Get("Content-Type")))
}

// AssertHTTPRequest asserts the value of an http.Request.
// Intended for use when testing incoming client requests
// See https://golang.org/pkg/net/http/httputil/#DumpRequest for more
func AssertHTTPRequest(t *testing.T, id string, r *http.Request) {
	body, err := httputil.DumpRequest(r, true)
	if err != nil {
		t.Fatal(err)
	}

	assertHTTP(t, id, body, contentTypeIsJSON(r.Header.Get("Content-Type")))
}

func assertHTTP(t *testing.T, id string, body []byte, isJSON bool) {
	config, err := getConfig()
	if err != nil {
		t.Fatal(err)
	}

	data := string(body)
	lines := strings.Split(strings.TrimSpace(data), "\n")

	// empty line identifies the end of the HTTP header
	for i, line := range lines {
		if line == "" {
			break
		}

		headerItem := strings.Split(line, ":")
		if def, ok := config.Defaults[headerItem[0]]; ok {
			lines[i] = fmt.Sprintf("%s: %s", headerItem[0], def)
		}
	}

	// If the response body is JSON, indent.
	if isJSON {
		jsonStr := lines[len(lines)-1]

		var jsonIface map[string]interface{}
		err = json.Unmarshal([]byte(jsonStr), &jsonIface)
		if err != nil {
			t.Fatal(err)
		}

		// Clean/update json based on config.
		if config != nil {
			for k, v := range config.Defaults {
				jsonIface = internal.UpdateKeyValuesInMap(k, v, jsonIface)
			}
		}

		out, err := json.MarshalIndent(jsonIface, "", "  ")
		if err != nil {
			t.Fatal(err)
		}
		lines[len(lines)-1] = string(out)
	}

	data = strings.Join(lines, "\n")
	createOrUpdateSnapshot(t, id, data)
}

func contentTypeIsJSON(contentType string) bool {
	contentTypeParts := strings.Split(contentType, ";")
	firstPart := contentTypeParts[0]

	isPlainJSON := firstPart == "application/json"
	if isPlainJSON {
		return isPlainJSON
	}

	isVendor := strings.HasPrefix(firstPart, "application/vnd.")

	isJSON := strings.HasSuffix(firstPart, "+json")

	return isVendor && isJSON
}

// AssertReader asserts the value of an io.Reader.
func AssertReader(t *testing.T, id string, r io.Reader) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}

	createOrUpdateSnapshot(t, id, string(data))
}

func createOrUpdateSnapshot(t *testing.T, id, data string) {
	var err error
	snapshot := getSnapshot(snapshotID(id))

	if snapshot == nil {
		if !args.shouldUpdate {
			t.Error(newSnapshotMessage(id, data))
			return
		}

		fmt.Printf("Creating snapshot `%s`\n", id)
		snapshot, err = createSnapshot(snapshotID(id), data)
		if err != nil {
			t.Fatal(err)
		}
		snapshot.evaluated = true
		return
	}

	snapshot.evaluated = true
	diff := compareResults(t, snapshot.value, strings.TrimSpace(data))
	if diff != "" {
		if snapshot != nil && args.shouldUpdate {
			fmt.Printf("Updating snapshot `%s`\n", id)
			_, err = createSnapshot(snapshotID(id), data)
			if err != nil {
				t.Fatal(err)
			}
			return
		}

		t.Error(didNotMatchMessage(id, diff))
		return
	}
}

func compareResults(t *testing.T, existing, new string) string {
	dmp := diffmatchpatch.New()
	dmp.PatchMargin = 20
	allDiffs := dmp.DiffMain(existing, new, false)
	nonEqualDiffs := []diffmatchpatch.Diff{}
	for _, diff := range allDiffs {
		if diff.Type != diffmatchpatch.DiffEqual {
			nonEqualDiffs = append(nonEqualDiffs, diff)
		}
	}

	if len(nonEqualDiffs) == 0 {
		return ""
	}

	return dmp.DiffPrettyText(allDiffs)
}

func didNotMatchMessage(id, diff string) string {
	msg := "\n\n## Existing snapshot does not match results...\n"
	msg += "## \"" + id + "\"\n\n"
	msg += diff
	msg += "\n\n"
	msg += "If this change was intentional, run tests again, $ go test -v -- -u\n"
	return msg
}

func newSnapshotMessage(id, body string) string {
	msg := "\n\n## New snapshot found...\n"
	msg += "## \"" + id + "\"\n\n"
	msg += body
	msg += "\n\n"
	msg += "To save, run tests again, $ go test -v -- -u\n"
	return msg
}
<commit_msg>Add safety check for nil config (#46)<commit_after>package abide

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"strings"
	"testing"

	"github.com/beme/abide/internal"
	"github.com/sergi/go-diff/diffmatchpatch"
)

// Assertable represents an object that can be asserted.
type Assertable interface {
	String() string
}

// Assert asserts the value of an object with implements Assertable.
func Assert(t *testing.T, id string, a Assertable) {
	data := a.String()
	createOrUpdateSnapshot(t, id, data)
}

// AssertHTTPResponse asserts the value of an http.Response.
func AssertHTTPResponse(t *testing.T, id string, w *http.Response) {
	body, err := httputil.DumpResponse(w, true)
	if err != nil {
		t.Fatal(err)
	}

	assertHTTP(t, id, body, contentTypeIsJSON(w.Header.Get("Content-Type")))
}

// AssertHTTPRequestOut asserts the value of an http.Request.
// Intended for use when testing outgoing client requests
// See https://golang.org/pkg/net/http/httputil/#DumpRequestOut for more
func AssertHTTPRequestOut(t *testing.T, id string, r *http.Request) {
	body, err := httputil.DumpRequestOut(r, true)
	if err != nil {
		t.Fatal(err)
	}

	assertHTTP(t, id, body, contentTypeIsJSON(r.Header.Get("Content-Type")))
}

// AssertHTTPRequest asserts the value of an http.Request.
// Intended for use when testing incoming client requests
// See https://golang.org/pkg/net/http/httputil/#DumpRequest for more
func AssertHTTPRequest(t *testing.T, id string, r *http.Request) {
	body, err := httputil.DumpRequest(r, true)
	if err != nil {
		t.Fatal(err)
	}

	assertHTTP(t, id, body, contentTypeIsJSON(r.Header.Get("Content-Type")))
}

func assertHTTP(t *testing.T, id string, body []byte, isJSON bool) {
	config, err := getConfig()
	if err != nil {
		t.Fatal(err)
	}

	data := string(body)
	lines := strings.Split(strings.TrimSpace(data), "\n")

	if config != nil {
		// empty line identifies the end of the HTTP header
		for i, line := range lines {
			if line == "" {
				break
			}

			headerItem := strings.Split(line, ":")
			if def, ok := config.Defaults[headerItem[0]]; ok {
				lines[i] = fmt.Sprintf("%s: %s", headerItem[0], def)
			}
		}
	}

	// If the response body is JSON, indent.
	if isJSON {
		jsonStr := lines[len(lines)-1]

		var jsonIface map[string]interface{}
		err = json.Unmarshal([]byte(jsonStr), &jsonIface)
		if err != nil {
			t.Fatal(err)
		}

		// Clean/update json based on config.
		if config != nil {
			for k, v := range config.Defaults {
				jsonIface = internal.UpdateKeyValuesInMap(k, v, jsonIface)
			}
		}

		out, err := json.MarshalIndent(jsonIface, "", "  ")
		if err != nil {
			t.Fatal(err)
		}
		lines[len(lines)-1] = string(out)
	}

	data = strings.Join(lines, "\n")
	createOrUpdateSnapshot(t, id, data)
}

func contentTypeIsJSON(contentType string) bool {
	contentTypeParts := strings.Split(contentType, ";")
	firstPart := contentTypeParts[0]

	isPlainJSON := firstPart == "application/json"
	if isPlainJSON {
		return isPlainJSON
	}

	isVendor := strings.HasPrefix(firstPart, "application/vnd.")

	isJSON := strings.HasSuffix(firstPart, "+json")

	return isVendor && isJSON
}

// AssertReader asserts the value of an io.Reader.
func AssertReader(t *testing.T, id string, r io.Reader) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}

	createOrUpdateSnapshot(t, id, string(data))
}

func createOrUpdateSnapshot(t *testing.T, id, data string) {
	var err error
	snapshot := getSnapshot(snapshotID(id))

	if snapshot == nil {
		if !args.shouldUpdate {
			t.Error(newSnapshotMessage(id, data))
			return
		}

		fmt.Printf("Creating snapshot `%s`\n", id)
		snapshot, err = createSnapshot(snapshotID(id), data)
		if err != nil {
			t.Fatal(err)
		}
		snapshot.evaluated = true
		return
	}

	snapshot.evaluated = true
	diff := compareResults(t, snapshot.value, strings.TrimSpace(data))
	if diff != "" {
		if snapshot != nil && args.shouldUpdate {
			fmt.Printf("Updating snapshot `%s`\n", id)
			_, err = createSnapshot(snapshotID(id), data)
			if err != nil {
				t.Fatal(err)
			}
			return
		}

		t.Error(didNotMatchMessage(id, diff))
		return
	}
}

func compareResults(t *testing.T, existing, new string) string {
	dmp := diffmatchpatch.New()
	dmp.PatchMargin = 20
	allDiffs := dmp.DiffMain(existing, new, false)
	nonEqualDiffs := []diffmatchpatch.Diff{}
	for _, diff := range allDiffs {
		if diff.Type != diffmatchpatch.DiffEqual {
			nonEqualDiffs = append(nonEqualDiffs, diff)
		}
	}

	if len(nonEqualDiffs) == 0 {
		return ""
	}

	return dmp.DiffPrettyText(allDiffs)
}

func didNotMatchMessage(id, diff string) string {
	msg := "\n\n## Existing snapshot does not match results...\n"
	msg += "## \"" + id + "\"\n\n"
	msg += diff
	msg += "\n\n"
	msg += "If this change was intentional, run tests again, $ go test -v -- -u\n"
	return msg
}

func newSnapshotMessage(id, body string) string {
	msg := "\n\n## New snapshot found...\n"
	msg += "## \"" + id + "\"\n\n"
	msg += body
	msg += "\n\n"
	msg += "To save, run tests again, $ go test -v -- -u\n"
	return msg
}
<|endoftext|>


<commit_before>// Copyright 2017 The go-darwin 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 hdiutil

import (
	"bytes"
	"fmt"
	"os/exec"
	"strconv"
)

// AttachFlag implements a hdiutil attach command flag interface.
type AttachFlag interface {
	attachFlag() []string
}

type rwType int

const (
	// Readonly force the resulting device to be read-only.
	readonly rwType = 1 << iota
	// Readwrite attempt to override the DiskImages framework's decision to attach a particular image read-only.
	// For example, -readwrite can be used to modify the HFS+ filesystem on a HFS+/ISO hybrid CD image.
	readwrite
)

func (rw rwType) attachFlag() string {
	switch rw {
	case readonly:
		return "-readonly"
	case readwrite:
		return "-readwrite"
	default:
		return ""
	}
}

type kernel bool

func (k kernel) attachFlag() []string { return boolNoFlag(bool(k), "kernel") }

type notRemovable bool

func (n notRemovable) attachFlag() []string { return boolFlag(bool(n), "notremovable") }

type mount string

func (m mount) attachFlag() []string { return stringFlag(string(m), "mount") }

type noMount bool

func (n noMount) attachFlag() []string { return boolFlag(bool(n), "nomount") }

// AttachMountRoot mount volumes on subdirectories of path instead of under /Volumes. path must exist.
// Full mount point paths must be less than MNAMELEN characters (increased from 90 to 1024 in Mac OS X 10.6).
type AttachMountRoot string

func (m AttachMountRoot) attachFlag() []string { return stringFlag(string(m), "mountroot") }

// AttachMountRandom like -mountroot, but mount point directory names are randomized with mkdtemp(3).
type AttachMountRandom string

func (m AttachMountRandom) attachFlag() []string { return stringFlag(string(m), "mountrandom") }

// AttachMountPoint assuming only one volume, mount it at path instead of in /Volumes.
// See fstab(5) for ways a system administrator can make particular volumes automatically mount in particular filesystem locations by editing the file /etc/fstab.
type AttachMountPoint string

func (m AttachMountPoint) attachFlag() []string { return stringFlag(string(m), "mountpoint") }

type noBrowse bool

func (n noBrowse) attachFlag() []string { return boolFlag(bool(n), "nobrowse") }

type owners string

const (
	ownersOn  owners = "on"
	ownersOff owners = "off"
)

func (o owners) attachFlag() []string { return stringFlag(string(o), "owners") }

// AttachDrivekey specify a key/value pair to be set on the device in the IOKit registry.
type AttachDrivekey [2]string

func (d AttachDrivekey) attachFlag() []string {
	return stringFlag(d[0]+"="+d[1], "drivekey")
}

// AttachSection attach a subsection of a disk image.
// subspec is any of <offset>, <first-last>, or <start,count> in 0-based sectors.
// Ranges are inclusive.
type AttachSection [2]int

func (s AttachSection) attachFlag() []string {
	var arg string
	for v := range s {
		arg = arg + strconv.Itoa(v)
	}
	return stringFlag(arg, "section")
}

type verify bool

func (v verify) attachFlag() []string { return boolNoFlag(bool(v), "verify") }

type ignoreBadChecksums bool

func (i ignoreBadChecksums) attachFlag() []string { return boolNoFlag(bool(i), "ignoreBadChecksums") }

type idme bool

func (i idme) attachFlag() []string { return boolNoFlag(bool(i), "idme") }

type idmeReveal bool

func (i idmeReveal) attachFlag() []string { return boolNoFlag(bool(i), "idmereveal") }

type idmeTrash bool

func (i idmeTrash) attachFlag() []string { return boolNoFlag(bool(i), "idmetrash") }

type autoOpen bool

func (a autoOpen) attachFlag() []string { return boolNoFlag(bool(a), "autoopen") }

type autoOpenRO bool

func (a autoOpenRO) attachFlag() []string { return boolNoFlag(bool(a), "autoopenro") }

type autoOpenRW bool

func (a autoOpenRW) attachFlag() []string { return boolNoFlag(bool(a), "autoopenrw") }

type autoFsck bool

func (a autoFsck) attachFlag() []string { return boolNoFlag(bool(a), "autofsck") }

const (
	// AttachReadonly force the resulting device to be read-only.
	AttachReadonly rwType = readonly

	// AttachReadWrite attempt to override the DiskImages framework's decision to attach a particular image read-only.
	//
	// For example, -readwrite can be used to modify the HFS+ filesystem on a HFS+/ISO hybrid CD image.
	AttachReadWrite rwType = readwrite

	// AttachKernel attempt to attach this image without a helper process; fail if unsupported.
	//
	// Only UDRW, UDRO, UDZO, ULFO, and UDSP images are supported in-kernel. Encryption and HTTP are supported by the kernel driver.
	AttachKernel kernel = true

	// AttachNoKernel attach with a helper process.  This is (again) the default as of Mac OS X 10.5.
	AttachNoKernel kernel = false

	// AttachNotRemovable prevent this image from being detached. Only root can use this option.
	//
	// A reboot is necessary to cleanly detach an image attached with -notremovable.
	AttachNotRemovable notRemovable = true

	// AttachMountRequired indicate to -mount required.
	AttachMountRequired mount = "required"

	// AttachMountOptional indicate to -mount optional.
	AttachMountOptional mount = "optional"

	// AttachMountSuppressed indicate to -mount suppressed.
	AttachMountSuppressed mount = "suppressed"

	// AttachNoMount identical to -mount suppressed
	AttachNoMount noMount = true

	// AttachNoBrowse render any volumes invisible in applications such as the macOS Finder.
	AttachNoBrowse noBrowse = true

	// AttachOwnersOn owners on any filesystems be honored.
	AttachOwnersOn owners = ownersOn

	// AttachOwnersOff owners on any filesystems be not honored.
	AttachOwnersOff owners = ownersOff

	// AttachVerify do verify the image.
	AttachVerify verify = true

	// AttachNoVerify do not verify the image.
	AttachNoVerify verify = false

	// AttachIgnoreBadChecksums bad checksums should be ignored.
	AttachIgnoreBadChecksums ignoreBadChecksums = true

	// AttachNoIgnoreBadChecksums bad checksums should be not ignored.
	AttachNoIgnoreBadChecksums ignoreBadChecksums = false

	// AttachIdme do perform IDME actions on IDME images.
	AttachIdme idme = true

	// AttachNoIdme do not perform IDME actions on IDME images.
	AttachNoIdme idme = false

	// AttachIdmeReveal do reveal (in the Finder) the results of IDME processing.
	AttachIdmeReveal idmeReveal = true

	// AttachNoIdmeReveal do not reveal (in the Finder) the results of IDME processing.
	AttachNoIdmeReveal idmeReveal = false

	// AttachIdmeTrash do put IDME images in the trash after processing.
	AttachIdmeTrash idmeTrash = true
	// AttachNoIdmeTrash do not put IDME images in the trash after processing.
	AttachNoIdmeTrash idmeTrash = false

	// AttachAutoOpen do not auto-open volumes (in the Finder) after attaching an image.
	AttachAutoOpen autoOpen = true

	// AttachNoAutoOpen do not auto-open volumes (in the Finder) after attaching an image.
	AttachNoAutoOpen autoOpen = false

	// AttachAutoOpenRO do auto-open read-only volumes.
	AttachAutoOpenRO autoOpenRO = true

	// AttachNoAutoOpenRO do not auto-open read-only volumes.
	AttachNoAutoOpenRO autoOpenRO = false

	// AttachAutoOpenRW do auto-open read/write volumes.
	AttachAutoOpenRW autoOpenRW = true

	// AttachNoAutoOpenRW do not auto-open read/write volumes.
	AttachNoAutoOpenRW autoOpenRW = false

	// AttachAutoFsck do force automatic file system checking before mounting a disk image.
	AttachAutoFsck autoFsck = true

	// AttachNoAutoFsck do not force automatic file system checking before mounting a disk image.
	AttachNoAutoFsck autoFsck = false
)

// Attach attach the image file. The returns device node path and error.
func Attach(image string, flags ...AttachFlag) (string, error) {
	cmd := exec.Command(hdiutilPath, "attach", image)

	if len(flags) != 0 {
		args := []string{}
		for _, f := range flags {
			args = append(args, f.attachFlag()...)
		}
		cmd.Args = append(cmd.Args, args...)
	}

	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("%v: %v", err, string(out))
	}

	return string(bytes.TrimSpace(out)), nil
}
<commit_msg>attach: fix Attach result to deviceNode only when set -mountpoint<commit_after>// Copyright 2017 The go-darwin 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 hdiutil

import (
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
)

// AttachFlag implements a hdiutil attach command flag interface.
type AttachFlag interface {
	attachFlag() []string
}

type rwType int

const (
	// Readonly force the resulting device to be read-only.
	readonly rwType = 1 << iota
	// Readwrite attempt to override the DiskImages framework's decision to attach a particular image read-only.
	// For example, -readwrite can be used to modify the HFS+ filesystem on a HFS+/ISO hybrid CD image.
	readwrite
)

func (rw rwType) attachFlag() string {
	switch rw {
	case readonly:
		return "-readonly"
	case readwrite:
		return "-readwrite"
	default:
		return ""
	}
}

type kernel bool

func (k kernel) attachFlag() []string { return boolNoFlag(bool(k), "kernel") }

type notRemovable bool

func (n notRemovable) attachFlag() []string { return boolFlag(bool(n), "notremovable") }

type mount string

func (m mount) attachFlag() []string { return stringFlag(string(m), "mount") }

type noMount bool

func (n noMount) attachFlag() []string { return boolFlag(bool(n), "nomount") }

// AttachMountRoot mount volumes on subdirectories of path instead of under /Volumes. path must exist.
// Full mount point paths must be less than MNAMELEN characters (increased from 90 to 1024 in Mac OS X 10.6).
type AttachMountRoot string

func (m AttachMountRoot) attachFlag() []string { return stringFlag(string(m), "mountroot") }

// AttachMountRandom like -mountroot, but mount point directory names are randomized with mkdtemp(3).
type AttachMountRandom string

func (m AttachMountRandom) attachFlag() []string { return stringFlag(string(m), "mountrandom") }

// AttachMountPoint assuming only one volume, mount it at path instead of in /Volumes.
// See fstab(5) for ways a system administrator can make particular volumes automatically mount in particular filesystem locations by editing the file /etc/fstab.
type AttachMountPoint string

func (m AttachMountPoint) attachFlag() []string { return stringFlag(string(m), "mountpoint") }

type noBrowse bool

func (n noBrowse) attachFlag() []string { return boolFlag(bool(n), "nobrowse") }

type owners string

const (
	ownersOn  owners = "on"
	ownersOff owners = "off"
)

func (o owners) attachFlag() []string { return stringFlag(string(o), "owners") }

// AttachDrivekey specify a key/value pair to be set on the device in the IOKit registry.
type AttachDrivekey [2]string

func (d AttachDrivekey) attachFlag() []string {
	return stringFlag(d[0]+"="+d[1], "drivekey")
}

// AttachSection attach a subsection of a disk image.
// subspec is any of <offset>, <first-last>, or <start,count> in 0-based sectors.
// Ranges are inclusive.
type AttachSection [2]int

func (s AttachSection) attachFlag() []string {
	var arg string
	for v := range s {
		arg = arg + strconv.Itoa(v)
	}
	return stringFlag(arg, "section")
}

type verify bool

func (v verify) attachFlag() []string { return boolNoFlag(bool(v), "verify") }

type ignoreBadChecksums bool

func (i ignoreBadChecksums) attachFlag() []string { return boolNoFlag(bool(i), "ignoreBadChecksums") }

type idme bool

func (i idme) attachFlag() []string { return boolNoFlag(bool(i), "idme") }

type idmeReveal bool

func (i idmeReveal) attachFlag() []string { return boolNoFlag(bool(i), "idmereveal") }

type idmeTrash bool

func (i idmeTrash) attachFlag() []string { return boolNoFlag(bool(i), "idmetrash") }

type autoOpen bool

func (a autoOpen) attachFlag() []string { return boolNoFlag(bool(a), "autoopen") }

type autoOpenRO bool

func (a autoOpenRO) attachFlag() []string { return boolNoFlag(bool(a), "autoopenro") }

type autoOpenRW bool

func (a autoOpenRW) attachFlag() []string { return boolNoFlag(bool(a), "autoopenrw") }

type autoFsck bool

func (a autoFsck) attachFlag() []string { return boolNoFlag(bool(a), "autofsck") }

const (
	// AttachReadonly force the resulting device to be read-only.
	AttachReadonly rwType = readonly

	// AttachReadWrite attempt to override the DiskImages framework's decision to attach a particular image read-only.
	//
	// For example, -readwrite can be used to modify the HFS+ filesystem on a HFS+/ISO hybrid CD image.
	AttachReadWrite rwType = readwrite

	// AttachKernel attempt to attach this image without a helper process; fail if unsupported.
	//
	// Only UDRW, UDRO, UDZO, ULFO, and UDSP images are supported in-kernel. Encryption and HTTP are supported by the kernel driver.
	AttachKernel kernel = true

	// AttachNoKernel attach with a helper process.  This is (again) the default as of Mac OS X 10.5.
	AttachNoKernel kernel = false

	// AttachNotRemovable prevent this image from being detached. Only root can use this option.
	//
	// A reboot is necessary to cleanly detach an image attached with -notremovable.
	AttachNotRemovable notRemovable = true

	// AttachMountRequired indicate to -mount required.
	AttachMountRequired mount = "required"

	// AttachMountOptional indicate to -mount optional.
	AttachMountOptional mount = "optional"

	// AttachMountSuppressed indicate to -mount suppressed.
	AttachMountSuppressed mount = "suppressed"

	// AttachNoMount identical to -mount suppressed
	AttachNoMount noMount = true

	// AttachNoBrowse render any volumes invisible in applications such as the macOS Finder.
	AttachNoBrowse noBrowse = true

	// AttachOwnersOn owners on any filesystems be honored.
	AttachOwnersOn owners = ownersOn

	// AttachOwnersOff owners on any filesystems be not honored.
	AttachOwnersOff owners = ownersOff

	// AttachVerify do verify the image.
	AttachVerify verify = true

	// AttachNoVerify do not verify the image.
	AttachNoVerify verify = false

	// AttachIgnoreBadChecksums bad checksums should be ignored.
	AttachIgnoreBadChecksums ignoreBadChecksums = true

	// AttachNoIgnoreBadChecksums bad checksums should be not ignored.
	AttachNoIgnoreBadChecksums ignoreBadChecksums = false

	// AttachIdme do perform IDME actions on IDME images.
	AttachIdme idme = true

	// AttachNoIdme do not perform IDME actions on IDME images.
	AttachNoIdme idme = false

	// AttachIdmeReveal do reveal (in the Finder) the results of IDME processing.
	AttachIdmeReveal idmeReveal = true

	// AttachNoIdmeReveal do not reveal (in the Finder) the results of IDME processing.
	AttachNoIdmeReveal idmeReveal = false

	// AttachIdmeTrash do put IDME images in the trash after processing.
	AttachIdmeTrash idmeTrash = true
	// AttachNoIdmeTrash do not put IDME images in the trash after processing.
	AttachNoIdmeTrash idmeTrash = false

	// AttachAutoOpen do not auto-open volumes (in the Finder) after attaching an image.
	AttachAutoOpen autoOpen = true

	// AttachNoAutoOpen do not auto-open volumes (in the Finder) after attaching an image.
	AttachNoAutoOpen autoOpen = false

	// AttachAutoOpenRO do auto-open read-only volumes.
	AttachAutoOpenRO autoOpenRO = true

	// AttachNoAutoOpenRO do not auto-open read-only volumes.
	AttachNoAutoOpenRO autoOpenRO = false

	// AttachAutoOpenRW do auto-open read/write volumes.
	AttachAutoOpenRW autoOpenRW = true

	// AttachNoAutoOpenRW do not auto-open read/write volumes.
	AttachNoAutoOpenRW autoOpenRW = false

	// AttachAutoFsck do force automatic file system checking before mounting a disk image.
	AttachAutoFsck autoFsck = true

	// AttachNoAutoFsck do not force automatic file system checking before mounting a disk image.
	AttachNoAutoFsck autoFsck = false
)

var attachRe = regexp.MustCompile(`/dev/disk[\d]+`)

// Attach attach the image file. The returns device node path and error.
func Attach(image string, flags ...AttachFlag) (string, error) {
	cmd := exec.Command(hdiutilPath, "attach", image)

	if len(flags) != 0 {
		args := []string{}
		for _, f := range flags {
			args = append(args, f.attachFlag()...)
		}
		cmd.Args = append(cmd.Args, args...)
	}

	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("%v: %v", err, string(out))
	}

	return string(attachRe.Find(out)), nil
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 sdb

import (
	"fmt"
)

func (d *domain) GetAttributes(
	item ItemName,
	constistentRead bool,
	attrNames []string) (attrs []Attribute, err error) {
	// Make sure the item name is legal.
	if item == "" {
		err = fmt.Errorf("Invalid item name; names must be non-empty.")
		return
	}

	if err = validateValue(string(item)); err != nil {
		err = fmt.Errorf("Invalid item name: %v", err)
		return
	}

	// Make sure attribute names are legal.
	for _, name := range attrNames {
		if name == "" {
			err = fmt.Errorf("Invalid attribute name; names must be non-empty.")
			return
		}

		if err = validateValue(name); err != nil {
			err = fmt.Errorf("Invalid attribute name: %v", err)
			return
		}
	}

	return nil, fmt.Errorf("TODO")
}
<commit_msg>Implemented more of GetAttributes.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 sdb

import (
	"fmt"
	"github.com/jacobsa/aws/exp/sdb/conn"
)

func (d *domain) GetAttributes(
	item ItemName,
	constistentRead bool,
	attrNames []string) (attrs []Attribute, err error) {
	// Make sure the item name is legal.
	if item == "" {
		err = fmt.Errorf("Invalid item name; names must be non-empty.")
		return
	}

	if err = validateValue(string(item)); err != nil {
		err = fmt.Errorf("Invalid item name: %v", err)
		return
	}

	// Make sure attribute names are legal.
	for _, name := range attrNames {
		if name == "" {
			err = fmt.Errorf("Invalid attribute name; names must be non-empty.")
			return
		}

		if err = validateValue(name); err != nil {
			err = fmt.Errorf("Invalid attribute name: %v", err)
			return
		}
	}

	// Create an appropriate request.
	//
	// Reference:
	//     http://goo.gl/MmaJA
	req := conn.Request{
		"DomainName": d.name,
		"ItemName": string(item),
	}

	if constistentRead {
		req["ConsistentRead"] = "true"
	}

	for i, name := range attrNames {
		req[fmt.Sprintf("AttributeName.%d", i)] = name
	}

	// Call the connection.
	if _, err = d.c.SendRequest(req); err != nil {
		err = fmt.Errorf("SendRequest: %v", err)
		return
	}

	return nil, fmt.Errorf("TODO")
}
<|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 os

import "syscall"

func isSymlink(stat *syscall.Stat_t) bool {
	return stat.Mode&syscall.S_IFMT == syscall.S_IFLNK
}

func fileInfoFromStat(name string, fi *FileInfo, lstat, stat *syscall.Stat_t) *FileInfo {
	fi.Dev = uint64(stat.Dev)
	fi.Ino = uint64(stat.Ino)
	fi.Nlink = uint64(stat.Nlink)
	fi.Mode = uint32(stat.Mode)
	fi.Uid = int(stat.Uid)
	fi.Gid = int(stat.Gid)
	fi.Rdev = uint64(stat.Rdev)
	fi.Size = uint64(stat.Size)
	fi.Blksize = int64(stat.Blksize)
	fi.Blocks = stat.Blocks
	fi.Atime_ns = syscall.TimespecToNsec(stat.Atimespec)
	fi.Mtime_ns = syscall.TimespecToNsec(stat.Mtimespec)
	fi.Ctime_ns = syscall.TimespecToNsec(stat.Ctimespec)
	for i := len(name) - 1; i >= 0; i-- {
		if name[i] == '/' {
			name = name[i+1:]
			break
		}
	}
	fi.Name = name
	if isSymlink(lstat) && !isSymlink(stat) {
		fi.FollowedSymlink = true
	}
	return fi
}
<commit_msg>os/stat_freebsd.go: maybe fix build for freebsd<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 os

import "syscall"

func isSymlink(stat *syscall.Stat_t) bool {
	return stat.Mode&syscall.S_IFMT == syscall.S_IFLNK
}

func fileInfoFromStat(name string, fi *FileInfo, lstat, stat *syscall.Stat_t) *FileInfo {
	fi.Dev = uint64(stat.Dev)
	fi.Ino = uint64(stat.Ino)
	fi.Nlink = uint64(stat.Nlink)
	fi.Mode = uint32(stat.Mode)
	fi.Uid = int(stat.Uid)
	fi.Gid = int(stat.Gid)
	fi.Rdev = uint64(stat.Rdev)
	fi.Size = int64(stat.Size)
	fi.Blksize = int64(stat.Blksize)
	fi.Blocks = stat.Blocks
	fi.Atime_ns = syscall.TimespecToNsec(stat.Atimespec)
	fi.Mtime_ns = syscall.TimespecToNsec(stat.Mtimespec)
	fi.Ctime_ns = syscall.TimespecToNsec(stat.Ctimespec)
	for i := len(name) - 1; i >= 0; i-- {
		if name[i] == '/' {
			name = name[i+1:]
			break
		}
	}
	fi.Name = name
	if isSymlink(lstat) && !isSymlink(stat) {
		fi.FollowedSymlink = true
	}
	return fi
}
<|endoftext|>


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

package runtime_test

import (
	"os"
	"runtime"
	"runtime/debug"
	"testing"
)

func TestGcSys(t *testing.T) {
	if os.Getenv("GOGC") == "off" {
		t.Fatalf("GOGC=off in environment; test cannot pass")
	}
	data := struct{ Short bool }{testing.Short()}
	got := executeTest(t, testGCSysSource, &data)
	want := "OK\n"
	if got != want {
		t.Fatalf("expected %q, but got %q", want, got)
	}
}

const testGCSysSource = `
package main

import (
	"fmt"
	"runtime"
)

func main() {
	runtime.GOMAXPROCS(1)
	memstats := new(runtime.MemStats)
	runtime.GC()
	runtime.ReadMemStats(memstats)
	sys := memstats.Sys

	runtime.MemProfileRate = 0 // disable profiler

	itercount := 1000000
{{if .Short}}
	itercount = 100000
{{end}}
	for i := 0; i < itercount; i++ {
		workthegc()
	}

	// Should only be using a few MB.
	// We allocated 100 MB or (if not short) 1 GB.
	runtime.ReadMemStats(memstats)
	if sys > memstats.Sys {
		sys = 0
	} else {
		sys = memstats.Sys - sys
	}
	if sys > 16<<20 {
		fmt.Printf("using too much memory: %d bytes\n", sys)
		return
	}
	fmt.Printf("OK\n")
}

func workthegc() []byte {
	return make([]byte, 1029)
}
`

func TestGcDeepNesting(t *testing.T) {
	type T [2][2][2][2][2][2][2][2][2][2]*int
	a := new(T)

	// Prevent the compiler from applying escape analysis.
	// This makes sure new(T) is allocated on heap, not on the stack.
	t.Logf("%p", a)

	a[0][0][0][0][0][0][0][0][0][0] = new(int)
	*a[0][0][0][0][0][0][0][0][0][0] = 13
	runtime.GC()
	if *a[0][0][0][0][0][0][0][0][0][0] != 13 {
		t.Fail()
	}
}

func TestGcHashmapIndirection(t *testing.T) {
	defer debug.SetGCPercent(debug.SetGCPercent(1))
	runtime.GC()
	type T struct {
		a [256]int
	}
	m := make(map[T]T)
	for i := 0; i < 2000; i++ {
		var a T
		a.a[0] = i
		m[a] = T{}
	}
}
<commit_msg>runtime: TestGcSys: if GOGC=off, skip instead of failing<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime_test

import (
	"os"
	"runtime"
	"runtime/debug"
	"testing"
)

func TestGcSys(t *testing.T) {
	if os.Getenv("GOGC") == "off" {
		t.Skip("skipping test; GOGC=off in environment")
	}
	data := struct{ Short bool }{testing.Short()}
	got := executeTest(t, testGCSysSource, &data)
	want := "OK\n"
	if got != want {
		t.Fatalf("expected %q, but got %q", want, got)
	}
}

const testGCSysSource = `
package main

import (
	"fmt"
	"runtime"
)

func main() {
	runtime.GOMAXPROCS(1)
	memstats := new(runtime.MemStats)
	runtime.GC()
	runtime.ReadMemStats(memstats)
	sys := memstats.Sys

	runtime.MemProfileRate = 0 // disable profiler

	itercount := 1000000
{{if .Short}}
	itercount = 100000
{{end}}
	for i := 0; i < itercount; i++ {
		workthegc()
	}

	// Should only be using a few MB.
	// We allocated 100 MB or (if not short) 1 GB.
	runtime.ReadMemStats(memstats)
	if sys > memstats.Sys {
		sys = 0
	} else {
		sys = memstats.Sys - sys
	}
	if sys > 16<<20 {
		fmt.Printf("using too much memory: %d bytes\n", sys)
		return
	}
	fmt.Printf("OK\n")
}

func workthegc() []byte {
	return make([]byte, 1029)
}
`

func TestGcDeepNesting(t *testing.T) {
	type T [2][2][2][2][2][2][2][2][2][2]*int
	a := new(T)

	// Prevent the compiler from applying escape analysis.
	// This makes sure new(T) is allocated on heap, not on the stack.
	t.Logf("%p", a)

	a[0][0][0][0][0][0][0][0][0][0] = new(int)
	*a[0][0][0][0][0][0][0][0][0][0] = 13
	runtime.GC()
	if *a[0][0][0][0][0][0][0][0][0][0] != 13 {
		t.Fail()
	}
}

func TestGcHashmapIndirection(t *testing.T) {
	defer debug.SetGCPercent(debug.SetGCPercent(1))
	runtime.GC()
	type T struct {
		a [256]int
	}
	m := make(map[T]T)
	for i := 0; i < 2000; i++ {
		var a T
		a.a[0] = i
		m[a] = T{}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"io"
	"regexp"
	"strings"
	"flag"
	"time"
	"compress/gzip"
)

const TIME_LAYOUT string = "[2006-01-02 15:04:05 MST]"
var job_regexp       *regexp.Regexp = regexp.MustCompile("^P[0-9]+(DJ|PW)[0-9]*")
var timestamp_regexp *regexp.Regexp = regexp.MustCompile("^(\\[[0-9-]+ [0-9:]+ UTC\\])")
var request_regexp   *regexp.Regexp = regexp.MustCompile("\\] (P[0-9]+[A-Za-z]+[0-9]+) ")

func main() {
	hide_jobs_flag := flag.Bool("hide_jobs", false, "Hide background jobs")
	hide_sql_flag  := flag.Bool("hide_sql", false, "Hide SQL statements")
	hide_ntlm_flag := flag.Bool("hide_ntlm", false, "Hide NTLM lines")
	full_flag      := flag.Bool("full", false, "Show the full request/job for each found line")
	after_str      := flag.String("after", "", "Show logs after this time (YYYY-MM-DD HH:II::SS")
	find_str       := flag.String("find", "", "Find lines matching this regexp")

	flag.Parse()
	args := flag.Args()

	// Time layouts must use the
	// reference time `Mon Jan 2 15:04:05 MST 2006` to show the
	// pattern with which to format/parse a given time/string

	time_after, e    := time.Parse(TIME_LAYOUT, fmt.Sprintf("[%s UTC]", *after_str))
	parse_time       := false

	if e != nil {
		if len(*after_str) > 0 {
			fmt.Println(fmt.Sprintf("Invalid time format \"%s\" - Must be YYYY-MM-DD HH::II::SS", *after_str))
			usage()
			os.Exit(2)
		}
	} else {
		parse_time = true
	}

	if len(args) < 1 {
		usage()
		os.Exit(2)
	}

	fmt.Println(fmt.Sprintf("Show full requests/jobs: %t", *full_flag))
	fmt.Println(fmt.Sprintf("Show background job lines: %t", !*hide_jobs_flag))
	fmt.Println(fmt.Sprintf("Show SQL lines: %t", !*hide_sql_flag))
	fmt.Println(fmt.Sprintf("Show NTLM lines: %t", !*hide_ntlm_flag))
	fmt.Println(fmt.Sprintf("Show lines after: %s", *after_str))

	filename := args[0]
	fmt.Println(fmt.Sprintf("Opening file: %s", filename))

	file := openFile(filename)
	defer file.Close()

	var reader io.Reader = file

	sql_regexp := regexp.MustCompile("(SQL \\()|(EXEC sp_executesql N)|( CACHE \\()")
	ntlm_regexp := regexp.MustCompile(" \\(NTLM\\) ")

	var unique_map map[string]bool;

	line_strexp := *find_str

	if line_regexp, err := regexp.Compile(line_strexp); *full_flag && len(line_strexp) > 0 && err == nil {
		if isGzip(filename) {
			// for some reason if you create a reader but don't use it,
			// an error is given when the output reader is created below
			parse_gz_reader := getGzipReader(file)
			defer parse_gz_reader.Close()

			reader = parse_gz_reader
		}

		line_count  := 0
		line_after  := !parse_time // if not parsing time, then all lines are valid
		request_ids := make([]string, 0)

		scanner := bufio.NewScanner(reader);

		for scanner.Scan() {
			line := scanner.Text();
			if line_regexp.MatchString(line) {

				if !line_after {
					if timestamp := extractTimestamp(line); len(timestamp) > 1 {
						if isAfterTime(timestamp, &time_after) {
							line_after = true
						}
					}
				}

				if line_after {
					if request_id := extractRequestId(line); len(request_id) > 1 {
						if !*hide_jobs_flag || !isJob(request_id) {
							request_ids = append(request_ids, request_id)
						}
					}
				}
			}

			if line_count++; line_count % 10000 == 0 {
				fmt.Print(fmt.Sprintf("Reading: %d\r", line_count))
			}
		}

		fmt.Println("") // empty line

		if err := scanner.Err(); err != nil {
			log.Fatal(err)
		}

		fmt.Println(fmt.Sprintf("Found %d lines matching \"%s\"", len(request_ids), line_strexp))
		unique_map = generateRequestIdMap(&request_ids)

		if len(unique_map) < 1 {
			fmt.Println(fmt.Sprintf("Found 0 request identifiers", line_strexp))
			os.Exit(2)
		}

		rewindFile(file)
	} else {
		fmt.Println("No matchers provided, skipping match phase")
	}

	if isGzip(filename) {
		output_gz_reader := getGzipReader(file)
		defer output_gz_reader.Close()

		reader = output_gz_reader
	}

	line_count := 0
	line_after := !parse_time // if not parsing time, then all lines are valid
	has_requests := len(unique_map) > 0

	line_regexp, err := regexp.Compile(line_strexp);
	has_matcher      := len(line_strexp) > 0 && err == nil

	output_scanner := bufio.NewScanner(reader);

	for output_scanner.Scan() {
		line := output_scanner.Text();

		output := false

		if !line_after {
			if line_count++; line_count % 10000 == 0 {
				fmt.Print(fmt.Sprintf("Reading: %d\r", line_count))
			}

			if timestamp := extractTimestamp(line); len(timestamp) > 1 {
				if isAfterTime(timestamp, &time_after) {
					fmt.Println("\n") // empty line
					line_after = true
				}
			}
		}

		if line_after {
			request_id := extractRequestId(line)

			if has_requests {
				if len(request_id) > 0 && unique_map[request_id] {
					if *hide_jobs_flag && isJob(request_id) {
						output = false
					} else {
						output = true
					}
				}
			} else if has_matcher {
				output = line_regexp.MatchString(line)
			} else {
				output = true
			}
		}

		if output {
			if *hide_sql_flag && sql_regexp.MatchString(line) {
				output = false
			} else if *hide_ntlm_flag && ntlm_regexp.MatchString(line) {
				output = false
			}
		}

		if output {
			fmt.Println(line)
		}
	}

	if err := output_scanner.Err(); err != nil {
		log.Fatal(err)
	}
}

func usage() {
	fmt.Println("Example: avm -find=\"username|computername\" \"/path/to/manager/log/production.log\"")

	flag.PrintDefaults()
}

func isAfterTime(timestamp string, time_after *time.Time) bool {
	if line_time, e := time.Parse(TIME_LAYOUT, timestamp); e != nil {
		fmt.Println("Got error %s", e)
		return false
	} else if line_time.Before(*time_after) {
		return false
	}

	return true
}

func isJob(request_id string) bool {
	return job_regexp.MatchString(request_id)
}

func extractTimestamp(line string) string {
	if timestamp_match := timestamp_regexp.FindStringSubmatch(line); len(timestamp_match) > 1 {
		return timestamp_match[1]
	} else {
		return ""
	}
}

func extractRequestId(line string) string {
	if request_match := request_regexp.FindStringSubmatch(line); len(request_match) > 1 {
		return request_match[1]
	} else {
		return ""
	}
}

func generateRequestIdMap(request_ids *[]string) map[string]bool {
	unique_map := make(map[string]bool, len(*request_ids))

	for _, x := range *request_ids {
		unique_map[x] = true
	}

	for k, _ := range unique_map {
		fmt.Println(fmt.Sprintf("Request ID: %s", k))
	}

	return unique_map
}

func openFile(filename string) *os.File {
	file, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	return file
}

func isGzip(filename string) bool {
	return strings.HasSuffix(filename, ".gz")
}

func getGzipReader(file *os.File) *gzip.Reader {
	gz_reader, err := gzip.NewReader(file)
	if err != nil {
	log.Fatal(err)
	}

	return gz_reader
}

func rewindFile(file *os.File) {
	file.Seek(0, 0)  // go back to the top (rewind)
}<commit_msg>Add the -neat flag as a shortcut for -hide_jobs, -hide_sql, -hide_ntlm<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"io"
	"regexp"
	"strings"
	"flag"
	"time"
	"compress/gzip"
)

const TIME_LAYOUT string = "[2006-01-02 15:04:05 MST]"
var job_regexp       *regexp.Regexp = regexp.MustCompile("^P[0-9]+(DJ|PW)[0-9]*")
var timestamp_regexp *regexp.Regexp = regexp.MustCompile("^(\\[[0-9-]+ [0-9:]+ UTC\\])")
var request_regexp   *regexp.Regexp = regexp.MustCompile("\\] (P[0-9]+[A-Za-z]+[0-9]+) ")

func main() {
	hide_jobs_flag := flag.Bool("hide_jobs", false, "Hide background jobs")
	hide_sql_flag  := flag.Bool("hide_sql", false, "Hide SQL statements")
	hide_ntlm_flag := flag.Bool("hide_ntlm", false, "Hide NTLM lines")
	full_flag      := flag.Bool("full", false, "Show the full request/job for each found line")
	neat_flag      := flag.Bool("neat", false, "Hide clutter - equivalent to -hide_jobs -hide_sql -hide_ntlm")
	after_str      := flag.String("after", "", "Show logs after this time (YYYY-MM-DD HH:II::SS")
	find_str       := flag.String("find", "", "Find lines matching this regexp")

	flag.Parse()
	args := flag.Args()

	// Time layouts must use the
	// reference time `Mon Jan 2 15:04:05 MST 2006` to show the
	// pattern with which to format/parse a given time/string

	time_after, e    := time.Parse(TIME_LAYOUT, fmt.Sprintf("[%s UTC]", *after_str))
	parse_time       := false

	if e != nil {
		if len(*after_str) > 0 {
			fmt.Println(fmt.Sprintf("Invalid time format \"%s\" - Must be YYYY-MM-DD HH::II::SS", *after_str))
			usage()
			os.Exit(2)
		}
	} else {
		parse_time = true
	}

	if len(args) < 1 {
		usage()
		os.Exit(2)
	}

	if *neat_flag {
		*hide_jobs_flag = true
		*hide_sql_flag = true
		*hide_ntlm_flag = true
	}

	fmt.Println(fmt.Sprintf("Show full requests/jobs: %t", *full_flag))
	fmt.Println(fmt.Sprintf("Show background job lines: %t", !*hide_jobs_flag))
	fmt.Println(fmt.Sprintf("Show SQL lines: %t", !*hide_sql_flag))
	fmt.Println(fmt.Sprintf("Show NTLM lines: %t", !*hide_ntlm_flag))
	fmt.Println(fmt.Sprintf("Show lines after: %s", *after_str))

	filename := args[0]
	fmt.Println(fmt.Sprintf("Opening file: %s", filename))

	file := openFile(filename)
	defer file.Close()

	var reader io.Reader = file

	sql_regexp := regexp.MustCompile("(SQL \\()|(EXEC sp_executesql N)|( CACHE \\()")
	ntlm_regexp := regexp.MustCompile(" \\(NTLM\\) ")

	var unique_map map[string]bool;

	line_strexp := *find_str

	if line_regexp, err := regexp.Compile(line_strexp); *full_flag && len(line_strexp) > 0 && err == nil {
		if isGzip(filename) {
			// for some reason if you create a reader but don't use it,
			// an error is given when the output reader is created below
			parse_gz_reader := getGzipReader(file)
			defer parse_gz_reader.Close()

			reader = parse_gz_reader
		}

		line_count  := 0
		line_after  := !parse_time // if not parsing time, then all lines are valid
		request_ids := make([]string, 0)

		scanner := bufio.NewScanner(reader);

		for scanner.Scan() {
			line := scanner.Text();
			if line_regexp.MatchString(line) {

				if !line_after {
					if timestamp := extractTimestamp(line); len(timestamp) > 1 {
						if isAfterTime(timestamp, &time_after) {
							line_after = true
						}
					}
				}

				if line_after {
					if request_id := extractRequestId(line); len(request_id) > 1 {
						if !*hide_jobs_flag || !isJob(request_id) {
							request_ids = append(request_ids, request_id)
						}
					}
				}
			}

			if line_count++; line_count % 10000 == 0 {
				fmt.Print(fmt.Sprintf("Reading: %d\r", line_count))
			}
		}

		fmt.Println("") // empty line

		if err := scanner.Err(); err != nil {
			log.Fatal(err)
		}

		fmt.Println(fmt.Sprintf("Found %d lines matching \"%s\"", len(request_ids), line_strexp))
		unique_map = generateRequestIdMap(&request_ids)

		if len(unique_map) < 1 {
			fmt.Println(fmt.Sprintf("Found 0 request identifiers", line_strexp))
			os.Exit(2)
		}

		rewindFile(file)
	} else {
		fmt.Println("No matchers provided, skipping match phase")
	}

	if isGzip(filename) {
		output_gz_reader := getGzipReader(file)
		defer output_gz_reader.Close()

		reader = output_gz_reader
	}

	line_count := 0
	line_after := !parse_time // if not parsing time, then all lines are valid
	has_requests := len(unique_map) > 0

	line_regexp, err := regexp.Compile(line_strexp);
	has_matcher      := len(line_strexp) > 0 && err == nil

	output_scanner := bufio.NewScanner(reader);

	for output_scanner.Scan() {
		line := output_scanner.Text();

		output := false

		if !line_after {
			if line_count++; line_count % 10000 == 0 {
				fmt.Print(fmt.Sprintf("Reading: %d\r", line_count))
			}

			if timestamp := extractTimestamp(line); len(timestamp) > 1 {
				if isAfterTime(timestamp, &time_after) {
					fmt.Println("\n") // empty line
					line_after = true
				}
			}
		}

		if line_after {
			request_id := extractRequestId(line)

			if has_requests {
				if len(request_id) > 0 && unique_map[request_id] {
					if *hide_jobs_flag && isJob(request_id) {
						output = false
					} else {
						output = true
					}
				}
			} else if has_matcher {
				output = line_regexp.MatchString(line)
			} else {
				output = true
			}
		}

		if output {
			if *hide_sql_flag && sql_regexp.MatchString(line) {
				output = false
			} else if *hide_ntlm_flag && ntlm_regexp.MatchString(line) {
				output = false
			}
		}

		if output {
			fmt.Println(line)
		}
	}

	if err := output_scanner.Err(); err != nil {
		log.Fatal(err)
	}
}

func usage() {
	fmt.Println("This tool can be used to extract the logs for specific requests from an AppVolumes Manager log")
	fmt.Println("Example:avmlog -after=\"2015-10-19 09:00:00\" -find \"apvuser2599\" -full -neat ~/Documents/scale.log.gz")

	flag.PrintDefaults()
}

func isAfterTime(timestamp string, time_after *time.Time) bool {
	if line_time, e := time.Parse(TIME_LAYOUT, timestamp); e != nil {
		fmt.Println("Got error %s", e)
		return false
	} else if line_time.Before(*time_after) {
		return false
	}

	return true
}

func isJob(request_id string) bool {
	return job_regexp.MatchString(request_id)
}

func extractTimestamp(line string) string {
	if timestamp_match := timestamp_regexp.FindStringSubmatch(line); len(timestamp_match) > 1 {
		return timestamp_match[1]
	} else {
		return ""
	}
}

func extractRequestId(line string) string {
	if request_match := request_regexp.FindStringSubmatch(line); len(request_match) > 1 {
		return request_match[1]
	} else {
		return ""
	}
}

func generateRequestIdMap(request_ids *[]string) map[string]bool {
	unique_map := make(map[string]bool, len(*request_ids))

	for _, x := range *request_ids {
		unique_map[x] = true
	}

	for k, _ := range unique_map {
		fmt.Println(fmt.Sprintf("Request ID: %s", k))
	}

	return unique_map
}

func openFile(filename string) *os.File {
	file, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	return file
}

func isGzip(filename string) bool {
	return strings.HasSuffix(filename, ".gz")
}

func getGzipReader(file *os.File) *gzip.Reader {
	gz_reader, err := gzip.NewReader(file)
	if err != nil {
	log.Fatal(err)
	}

	return gz_reader
}

func rewindFile(file *os.File) {
	file.Seek(0, 0)  // go back to the top (rewind)
}<|endoftext|>