<commit_before>package main

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/mattn/go-zglob"

	"github.com/geek1011/kepubify/kepub"

	cli "gopkg.in/urfave/cli.v1"
)

var version = "dev"

// exists checks whether a path exists
func exists(path string) bool {
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		return true
	}
	return false
}

// isDir checks if a exists and is a dir
func isDir(path string) bool {
	if fi, err := os.Stat(path); err == nil && fi.IsDir() {
		return true
	}
	return false
}

func convert(c *cli.Context) error {
	if len(c.Args()) < 1 || len(c.Args()) > 2 {
		return fmt.Errorf("Invalid number of arguments. Usage: kepubify EPUB_INPUT_PATH [KEPUB_OUTPUT_PATH]")
	}

	infile := c.Args().Get(0)
	if infile == "" {
		return fmt.Errorf("Input path must not be blank.")
	}

	infile, err := filepath.Abs(infile)
	if err != nil {
		return fmt.Errorf("Error resolving input path: %s.", err)
	}

	if !exists(infile) {
		return fmt.Errorf("Input file or directory does not exist.")
	}

	if isDir(infile) {
		if len(c.Args()) != 2 {
			return fmt.Errorf("Because input is a dir, a second argument must be supplied with a nonexistent dir for the conversion output.")
		}

		outdir := c.Args().Get(1)
		if exists(outdir) {
			return fmt.Errorf("Because input is a dir, a second argument must be supplied with a nonexistent dir for the conversion output.")
		}

		outdir, err = filepath.Abs(outdir)
		if err != nil {
			return fmt.Errorf("Error resolving output dir path: %s.", err)
		}

		err := os.Mkdir(outdir, os.ModePerm)
		if err != nil {
			return fmt.Errorf("Error creating output dir: %s.", err)
		}

		lst, err := zglob.Glob(filepath.Join(infile, "**", "*.epub"))
		if err != nil {
			return fmt.Errorf("Error searching for epubs in input dir: %s.", err)
		}

		epubs := []string{}
		for _, f := range lst {
			if !strings.HasSuffix(f, ".kepub.epub") {
				epubs = append(epubs, f)
			}
		}

		fmt.Printf("%v books found\n", len(epubs))

		errs := map[string]error{}
		for i, epub := range epubs {
			rel, err := filepath.Rel(infile, epub)
			if err != nil {
				fmt.Printf("[%v/%v] Error resolving relative path of %s: %v\n", i+1, len(epubs), epub, err)
				errs[epub] = err
				continue
			}

			err = os.MkdirAll(filepath.Join(outdir, filepath.Dir(rel)), os.ModePerm)
			if err != nil {
				fmt.Printf("[%v/%v] Error creating output dir for %s: %v\n", i+1, len(epubs), epub, err)
				errs[rel] = err
				continue
			}

			outfile := fmt.Sprintf("%s.kepub.epub", filepath.Join(outdir, strings.Replace(rel, ".epub", "", -1)))
			fmt.Printf("[%v/%v] Converting %s\n", i+1, len(epubs), rel)

			err = kepub.Kepubify(epub, outfile, false)
			if err != nil {
				fmt.Printf("[%v/%v] Error converting %s: %v\n", i+1, len(epubs), rel, err)
				errs[rel] = err
				continue
			}
		}

		fmt.Printf("\nSucessfully converted %v of %v ebooks\n", len(epubs)-len(errs), len(epubs))
		if len(errs) > 0 {
			fmt.Printf("Errors:\n")
			for epub, err := range errs {
				fmt.Printf("%s: %v\n", epub, err)
			}
		}
	} else {
		if filepath.Ext(infile) != ".epub" {
			return fmt.Errorf("Input file must have the extension \".epub\".")
		}

		outfile := fmt.Sprintf("%s.kepub.epub", strings.Replace(filepath.Base(infile), ".epub", "", -1))
		if len(c.Args()) == 2 {
			if exists(c.Args().Get(1)) {
				if isDir(c.Args().Get(1)) {
					outfile = path.Join(c.Args().Get(1), outfile)
				} else {
					return fmt.Errorf("Output path must be a nonexistent file ending with .kepub.epub, or be an existing directory")
				}
			} else {
				if strings.HasSuffix(c.Args().Get(1), ".kepub.epub") {
					outfile = c.Args().Get(1)
				} else {
					return fmt.Errorf("Output path must be a nonexistent file ending with .kepub.epub, or be an existing directory")
				}
			}
		}

		outfile, err = filepath.Abs(outfile)
		if err != nil {
			return fmt.Errorf("Error resolving output file path: %s.", err)
		}

		fmt.Printf("Input file: %s\n", infile)
		fmt.Printf("Output file: %s\n", outfile)
		fmt.Println()

		err = kepub.Kepubify(infile, outfile, true)
		if err != nil {
			return fmt.Errorf("Error converting epub to kepub: %s.", err)
		}

		fmt.Printf("Successfully converted \"%s\" to a kepub.\nYou can find the converted file at \"%s\"\n", infile, outfile)
	}
	return nil
}

func main() {
	app := cli.NewApp()

	app.Name = "kepubify"
	app.Usage = "Convert epub to kepub"
	app.Description = "Convert your ePubs into kepubs, with a easy-to-use command-line tool."
	app.Version = version

	app.ArgsUsage = "EPUB_INPUT_PATH [KEPUB_OUTPUT_PATH]"
	app.Action = func(c *cli.Context) error {
		err := convert(c)
		if err != nil {
			fmt.Println(err)
		}

		if runtime.GOOS == "windows" && len(c.Args()) == 1 {
			time.Sleep(1 * time.Second)
		}

		return err
	}

	app.Run(os.Args)
}
<commit_msg>Improved help text<commit_after>package main

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/mattn/go-zglob"

	"github.com/geek1011/kepubify/kepub"

	cli "gopkg.in/urfave/cli.v1"
)

var version = "dev"

// exists checks whether a path exists
func exists(path string) bool {
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		return true
	}
	return false
}

// isDir checks if a exists and is a dir
func isDir(path string) bool {
	if fi, err := os.Stat(path); err == nil && fi.IsDir() {
		return true
	}
	return false
}

func convert(c *cli.Context) error {
	if len(c.Args()) < 1 || len(c.Args()) > 2 {
		return fmt.Errorf("Invalid number of arguments.\n\n%s", helpText)
	}

	infile := c.Args().Get(0)
	if infile == "" {
		return fmt.Errorf("Input path must not be blank.")
	}

	infile, err := filepath.Abs(infile)
	if err != nil {
		return fmt.Errorf("Error resolving input path: %s.", err)
	}

	if !exists(infile) {
		return fmt.Errorf("Input file or directory does not exist.")
	}

	if isDir(infile) {
		if len(c.Args()) != 2 {
			return fmt.Errorf("Because input is a dir, a second argument must be supplied with a nonexistent dir for the conversion output.")
		}

		outdir := c.Args().Get(1)
		if exists(outdir) {
			return fmt.Errorf("Because input is a dir, a second argument must be supplied with a nonexistent dir for the conversion output.")
		}

		outdir, err = filepath.Abs(outdir)
		if err != nil {
			return fmt.Errorf("Error resolving output dir path: %s.", err)
		}

		err := os.Mkdir(outdir, os.ModePerm)
		if err != nil {
			return fmt.Errorf("Error creating output dir: %s.", err)
		}

		lst, err := zglob.Glob(filepath.Join(infile, "**", "*.epub"))
		if err != nil {
			return fmt.Errorf("Error searching for epubs in input dir: %s.", err)
		}

		epubs := []string{}
		for _, f := range lst {
			if !strings.HasSuffix(f, ".kepub.epub") {
				epubs = append(epubs, f)
			}
		}

		fmt.Printf("%v books found\n", len(epubs))

		errs := map[string]error{}
		for i, epub := range epubs {
			rel, err := filepath.Rel(infile, epub)
			if err != nil {
				fmt.Printf("[%v/%v] Error resolving relative path of %s: %v\n", i+1, len(epubs), epub, err)
				errs[epub] = err
				continue
			}

			err = os.MkdirAll(filepath.Join(outdir, filepath.Dir(rel)), os.ModePerm)
			if err != nil {
				fmt.Printf("[%v/%v] Error creating output dir for %s: %v\n", i+1, len(epubs), epub, err)
				errs[rel] = err
				continue
			}

			outfile := fmt.Sprintf("%s.kepub.epub", filepath.Join(outdir, strings.Replace(rel, ".epub", "", -1)))
			fmt.Printf("[%v/%v] Converting %s\n", i+1, len(epubs), rel)

			err = kepub.Kepubify(epub, outfile, false)
			if err != nil {
				fmt.Printf("[%v/%v] Error converting %s: %v\n", i+1, len(epubs), rel, err)
				errs[rel] = err
				continue
			}
		}

		fmt.Printf("\nSucessfully converted %v of %v ebooks\n", len(epubs)-len(errs), len(epubs))
		if len(errs) > 0 {
			fmt.Printf("Errors:\n")
			for epub, err := range errs {
				fmt.Printf("%s: %v\n", epub, err)
			}
		}
	} else {
		if filepath.Ext(infile) != ".epub" {
			return fmt.Errorf("Input file must have the extension \".epub\".")
		}

		outfile := fmt.Sprintf("%s.kepub.epub", strings.Replace(filepath.Base(infile), ".epub", "", -1))
		if len(c.Args()) == 2 {
			if exists(c.Args().Get(1)) {
				if isDir(c.Args().Get(1)) {
					outfile = path.Join(c.Args().Get(1), outfile)
				} else {
					return fmt.Errorf("Output path must be a nonexistent file ending with .kepub.epub, or be an existing directory")
				}
			} else {
				if strings.HasSuffix(c.Args().Get(1), ".kepub.epub") {
					outfile = c.Args().Get(1)
				} else {
					return fmt.Errorf("Output path must be a nonexistent file ending with .kepub.epub, or be an existing directory")
				}
			}
		}

		outfile, err = filepath.Abs(outfile)
		if err != nil {
			return fmt.Errorf("Error resolving output file path: %s.", err)
		}

		fmt.Printf("Input file: %s\n", infile)
		fmt.Printf("Output file: %s\n", outfile)
		fmt.Println()

		err = kepub.Kepubify(infile, outfile, true)
		if err != nil {
			return fmt.Errorf("Error converting epub to kepub: %s.", err)
		}

		fmt.Printf("Successfully converted \"%s\" to a kepub.\nYou can find the converted file at \"%s\"\n", infile, outfile)
	}
	return nil
}

var helpText = `USAGE: kepubify INPUT_PATH [OUTPUT_PATH]

VERSION: {{.Version}}

INPUT_PATH:
   The input file or directory. If it is a directory,
   OUTPUT_PATH must be a nonexistent directory.

OUTPUT_PATH:
   The path to place the converted ebook(s). Can only
   be a directory if INPUT_PATH is a directory.

   By default, this is the basename of the input file,
   with the extension .kepub.epub

The full documentation is available at:
   https://geek1011.github.io/kepubify
`

func main() {
	helpText = strings.Replace(helpText, "{{.Version}}", version, 1)

	app := cli.NewApp()

	app.Name = "kepubify"
	app.Usage = "An epub to kepub converter"
	app.Description = "Convert your ePubs into kepubs, with a easy-to-use command-line tool."
	app.Version = version

	app.ArgsUsage = "EPUB_INPUT_PATH [KEPUB_OUTPUT_PATH]"
	cli.AppHelpTemplate = helpText

	app.Action = func(c *cli.Context) error {
		err := convert(c)
		if err != nil {
			fmt.Println(err)
		}

		if runtime.GOOS == "windows" && len(c.Args()) == 1 {
			time.Sleep(1 * time.Second)
		}

		return err
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>// Copyright 2012 Google, Inc. All rights reserved.
// Copyright 2009-2011 Andreas Krennmair. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package layers

import (
	"code.google.com/p/gopacket"
	"encoding/binary"
	"fmt"
	"net"
	"strings"
)

type IPv4Flag uint8

const (
	IPv4EvilBit       IPv4Flag = 1 << 2 // http://tools.ietf.org/html/rfc3514 ;)
	IPv4DontFragment  IPv4Flag = 1 << 1
	IPv4MoreFragments IPv4Flag = 1 << 0
)

func (f IPv4Flag) String() string {
	var s []string
	if f&IPv4EvilBit != 0 {
		s = append(s, "Evil")
	}
	if f&IPv4DontFragment != 0 {
		s = append(s, "DF")
	}
	if f&IPv4MoreFragments != 0 {
		s = append(s, "MF")
	}
	return strings.Join(s, "|")
}

// IPv4 is the header of an IP packet.
type IPv4 struct {
	BaseLayer
	Version    uint8
	IHL        uint8
	TOS        uint8
	Length     uint16
	Id         uint16
	Flags      IPv4Flag
	FragOffset uint16
	TTL        uint8
	Protocol   IPProtocol
	Checksum   uint16
	SrcIP      net.IP
	DstIP      net.IP
	Options    []IPv4Option
	Padding    []byte
}

// LayerType returns LayerTypeIPv4
func (i *IPv4) LayerType() gopacket.LayerType { return LayerTypeIPv4 }
func (i *IPv4) NetworkFlow() gopacket.Flow {
	return gopacket.NewFlow(EndpointIPv4, i.SrcIP, i.DstIP)
}

type IPv4Option struct {
	OptionType   uint8
	OptionLength uint8
	OptionData   []byte
}

func (i IPv4Option) String() string {
	return fmt.Sprintf("IPv4Option(%v:%v)", i.OptionType, i.OptionData)
}

// SerializeTo writes the serialized form of this layer into the
// SerializationBuffer, implementing gopacket.SerializableLayer.
func (ip *IPv4) SerializeTo(b gopacket.SerializeBuffer, opts gopacket.SerializeOptions) error {
	if len(ip.Options) > 0 {
		return fmt.Errorf("cannot currently serialize IPv4 options")
	}
	bytes, err := b.PrependBytes(20)
	if err != nil {
		return err
	}
	if opts.FixLengths {
		ip.IHL = 5 // Fix when we add support for options.
		ip.Length = uint16(len(b.Bytes()))
	}
	bytes[0] = (ip.Version << 4) | ip.IHL
	bytes[1] = ip.TOS
	binary.BigEndian.PutUint16(bytes[2:], ip.Length)
	binary.BigEndian.PutUint16(bytes[4:], ip.Id)
	binary.BigEndian.PutUint16(bytes[6:], ip.flagsfrags())
	bytes[8] = ip.TTL
	bytes[9] = byte(ip.Protocol)
	if len(ip.SrcIP) != 4 {
		return fmt.Errorf("invalid src IP %v", ip.SrcIP)
	}
	if len(ip.DstIP) != 4 {
		return fmt.Errorf("invalid dst IP %v", ip.DstIP)
	}
	copy(bytes[12:16], ip.SrcIP)
	copy(bytes[16:20], ip.DstIP)
	if opts.ComputeChecksums {
		// Clear checksum bytes
		bytes[10] = 0
		bytes[11] = 0
		// Compute checksum
		var csum uint32
		for i := 0; i < len(bytes); i += 2 {
			csum += uint32(bytes[i]) << 8
			csum += uint32(bytes[i+1])
		}
		ip.Checksum = ^uint16((csum >> 16) + csum)
	}
	binary.BigEndian.PutUint16(bytes[10:], ip.Checksum)
	return nil
}

func (ip *IPv4) flagsfrags() (ff uint16) {
	ff |= uint16(ip.Flags) << 13
	ff |= ip.FragOffset
	return
}

// DecodeFromBytes decodes the given bytes into this layer.
func (ip *IPv4) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	flagsfrags := binary.BigEndian.Uint16(data[6:8])

	ip.Version = uint8(data[0]) >> 4
	ip.IHL = uint8(data[0]) & 0x0F
	ip.TOS = data[1]
	ip.Length = binary.BigEndian.Uint16(data[2:4])
	ip.Id = binary.BigEndian.Uint16(data[4:6])
	ip.Flags = IPv4Flag(flagsfrags >> 13)
	ip.FragOffset = flagsfrags & 0x1FFF
	ip.TTL = data[8]
	ip.Protocol = IPProtocol(data[9])
	ip.Checksum = binary.BigEndian.Uint16(data[10:12])
	ip.SrcIP = data[12:16]
	ip.DstIP = data[16:20]
	// Set up an initial guess for contents/payload... we'll reset these soon.
	ip.BaseLayer = BaseLayer{Contents: data}

	if ip.Length < 20 {
		return fmt.Errorf("Invalid (too small) IP length (%d < 20)", ip.Length)
	} else if ip.IHL < 5 {
		return fmt.Errorf("Invalid (too small) IP header length (%d < 5)", ip.IHL)
	} else if int(ip.IHL*4) > int(ip.Length) {
		return fmt.Errorf("Invalid IP header length > IP length (%d > %d)", ip.IHL, ip.Length)
	}
	if cmp := len(data) - int(ip.Length); cmp > 0 {
		data = data[:ip.Length]
	} else if cmp < 0 {
		df.SetTruncated()
		if int(ip.IHL)*4 > len(data) {
			return fmt.Errorf("Not all IP header bytes available")
		}
	}
	ip.Contents = data[:ip.IHL*4]
	ip.Payload = data[ip.IHL*4:]
	// From here on, data contains the header options.
	data = data[20 : ip.IHL*4]
	// Pull out IP options
	for len(data) > 0 {
		if ip.Options == nil {
			// Pre-allocate to avoid growing the slice too much.
			ip.Options = make([]IPv4Option, 0, 4)
		}
		opt := IPv4Option{OptionType: data[0]}
		switch opt.OptionType {
		case 0: // End of options
			opt.OptionLength = 1
			ip.Options = append(ip.Options, opt)
			ip.Padding = data[1:]
			break
		case 1: // 1 byte padding
			opt.OptionLength = 1
		default:
			opt.OptionLength = data[1]
			opt.OptionData = data[2:opt.OptionLength]
		}
		if len(data) >= int(opt.OptionLength) {
			data = data[opt.OptionLength:]
		} else {
			return fmt.Errorf("IP option length exceeds remaining IP header size, option type %v length %v", opt.OptionType, opt.OptionLength)
		}
		ip.Options = append(ip.Options, opt)
	}
	return nil
}

func (i *IPv4) CanDecode() gopacket.LayerClass {
	return LayerTypeIPv4
}

func (i *IPv4) NextLayerType() gopacket.LayerType {
	if i.Flags&IPv4MoreFragments != 0 || i.FragOffset != 0 {
		return gopacket.LayerTypeFragment
	}
	return i.Protocol.LayerType()
}

func decodeIPv4(data []byte, p gopacket.PacketBuilder) error {
	ip := &IPv4{}
	err := ip.DecodeFromBytes(data, p)
	p.AddLayer(ip)
	p.SetNetworkLayer(ip)
	if err != nil {
		return err
	}
	return p.NextDecoder(ip.NextLayerType())
}
<commit_msg>fix(ipv4): More fragments flags was incorrectly reported as Evil (flags were reverded)<commit_after>// Copyright 2012 Google, Inc. All rights reserved.
// Copyright 2009-2011 Andreas Krennmair. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package layers

import (
	"code.google.com/p/gopacket"
	"encoding/binary"
	"fmt"
	"net"
	"strings"
)

type IPv4Flag uint8

const (
	IPv4EvilBit       IPv4Flag = 1 << 2 // http://tools.ietf.org/html/rfc3514 ;)
	IPv4DontFragment  IPv4Flag = 1 << 1
	IPv4MoreFragments IPv4Flag = 1 << 0
)

func (f IPv4Flag) String() string {
	var s []string

	if f&IPv4EvilBit != 0 {
		s = append(s, "Evil")
	}
	if f&IPv4DontFragment != 0 {
		s = append(s, "DF")
	}
	if f&IPv4MoreFragments != 0 {
		s = append(s, "MF")
	}
	return strings.Join(s, "|")
}

// IPv4 is the header of an IP packet.
type IPv4 struct {
	BaseLayer
	Version    uint8
	IHL        uint8
	TOS        uint8
	Length     uint16
	Id         uint16
	Flags      IPv4Flag
	FragOffset uint16
	TTL        uint8
	Protocol   IPProtocol
	Checksum   uint16
	SrcIP      net.IP
	DstIP      net.IP
	Options    []IPv4Option
	Padding    []byte
}

// LayerType returns LayerTypeIPv4
func (i *IPv4) LayerType() gopacket.LayerType { return LayerTypeIPv4 }
func (i *IPv4) NetworkFlow() gopacket.Flow {
	return gopacket.NewFlow(EndpointIPv4, i.SrcIP, i.DstIP)
}

type IPv4Option struct {
	OptionType   uint8
	OptionLength uint8
	OptionData   []byte
}

func (i IPv4Option) String() string {
	return fmt.Sprintf("IPv4Option(%v:%v)", i.OptionType, i.OptionData)
}

// SerializeTo writes the serialized form of this layer into the
// SerializationBuffer, implementing gopacket.SerializableLayer.
func (ip *IPv4) SerializeTo(b gopacket.SerializeBuffer, opts gopacket.SerializeOptions) error {
	if len(ip.Options) > 0 {
		return fmt.Errorf("cannot currently serialize IPv4 options")
	}
	bytes, err := b.PrependBytes(20)
	if err != nil {
		return err
	}
	if opts.FixLengths {
		ip.IHL = 5 // Fix when we add support for options.
		ip.Length = uint16(len(b.Bytes()))
	}
	bytes[0] = (ip.Version << 4) | ip.IHL
	bytes[1] = ip.TOS
	binary.BigEndian.PutUint16(bytes[2:], ip.Length)
	binary.BigEndian.PutUint16(bytes[4:], ip.Id)
	binary.BigEndian.PutUint16(bytes[6:], ip.flagsfrags())
	bytes[8] = ip.TTL
	bytes[9] = byte(ip.Protocol)
	if len(ip.SrcIP) != 4 {
		return fmt.Errorf("invalid src IP %v", ip.SrcIP)
	}
	if len(ip.DstIP) != 4 {
		return fmt.Errorf("invalid dst IP %v", ip.DstIP)
	}
	copy(bytes[12:16], ip.SrcIP)
	copy(bytes[16:20], ip.DstIP)
	if opts.ComputeChecksums {
		// Clear checksum bytes
		bytes[10] = 0
		bytes[11] = 0
		// Compute checksum
		var csum uint32
		for i := 0; i < len(bytes); i += 2 {
			csum += uint32(bytes[i]) << 8
			csum += uint32(bytes[i+1])
		}
		ip.Checksum = ^uint16((csum >> 16) + csum)
	}
	binary.BigEndian.PutUint16(bytes[10:], ip.Checksum)
	return nil
}

func (ip *IPv4) flagsfrags() (ff uint16) {
	ff |= uint16(ip.Flags) << 13
	ff |= ip.FragOffset
	return
}

// DecodeFromBytes decodes the given bytes into this layer.
func (ip *IPv4) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	flagsfrags := binary.BigEndian.Uint16(data[6:8])

	ip.Version = uint8(data[0]) >> 4
	ip.IHL = uint8(data[0]) & 0x0F
	ip.TOS = data[1]
	ip.Length = binary.BigEndian.Uint16(data[2:4])
	ip.Id = binary.BigEndian.Uint16(data[4:6])
	ip.Flags = IPv4Flag(flagsfrags >> 13)
	ip.FragOffset = flagsfrags & 0x1FFF
	ip.TTL = data[8]
	ip.Protocol = IPProtocol(data[9])
	ip.Checksum = binary.BigEndian.Uint16(data[10:12])
	ip.SrcIP = data[12:16]
	ip.DstIP = data[16:20]
	// Set up an initial guess for contents/payload... we'll reset these soon.
	ip.BaseLayer = BaseLayer{Contents: data}

	if ip.Length < 20 {
		return fmt.Errorf("Invalid (too small) IP length (%d < 20)", ip.Length)
	} else if ip.IHL < 5 {
		return fmt.Errorf("Invalid (too small) IP header length (%d < 5)", ip.IHL)
	} else if int(ip.IHL*4) > int(ip.Length) {
		return fmt.Errorf("Invalid IP header length > IP length (%d > %d)", ip.IHL, ip.Length)
	}
	if cmp := len(data) - int(ip.Length); cmp > 0 {
		data = data[:ip.Length]
	} else if cmp < 0 {
		df.SetTruncated()
		if int(ip.IHL)*4 > len(data) {
			return fmt.Errorf("Not all IP header bytes available")
		}
	}
	ip.Contents = data[:ip.IHL*4]
	ip.Payload = data[ip.IHL*4:]
	// From here on, data contains the header options.
	data = data[20 : ip.IHL*4]
	// Pull out IP options
	for len(data) > 0 {
		if ip.Options == nil {
			// Pre-allocate to avoid growing the slice too much.
			ip.Options = make([]IPv4Option, 0, 4)
		}
		opt := IPv4Option{OptionType: data[0]}
		switch opt.OptionType {
		case 0: // End of options
			opt.OptionLength = 1
			ip.Options = append(ip.Options, opt)
			ip.Padding = data[1:]
			break
		case 1: // 1 byte padding
			opt.OptionLength = 1
		default:
			opt.OptionLength = data[1]
			opt.OptionData = data[2:opt.OptionLength]
		}
		if len(data) >= int(opt.OptionLength) {
			data = data[opt.OptionLength:]
		} else {
			return fmt.Errorf("IP option length exceeds remaining IP header size, option type %v length %v", opt.OptionType, opt.OptionLength)
		}
		ip.Options = append(ip.Options, opt)
	}
	return nil
}

func (i *IPv4) CanDecode() gopacket.LayerClass {
	return LayerTypeIPv4
}

func (i *IPv4) NextLayerType() gopacket.LayerType {
	if i.Flags&IPv4MoreFragments != 0 || i.FragOffset != 0 {
		return gopacket.LayerTypeFragment
	}
	return i.Protocol.LayerType()
}

func decodeIPv4(data []byte, p gopacket.PacketBuilder) error {
	ip := &IPv4{}
	err := ip.DecodeFromBytes(data, p)
	p.AddLayer(ip)
	p.SetNetworkLayer(ip)
	if err != nil {
		return err
	}
	return p.NextDecoder(ip.NextLayerType())
}
<|endoftext|>


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

package containers

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/projectcalico/felix/fv/utils"
	"github.com/projectcalico/libcalico-go/lib/set"
	log "github.com/sirupsen/logrus"
)

type Container struct {
	Name     string
	IP       string
	Hostname string
	Stop     func()
	stopped  bool
	binaries set.Set
}

var containerIdx = 0

var runningContainers = []*Container{}

func Run(namePrefix string, args ...string) (c *Container) {

	// Build unique container name and struct.
	containerIdx++
	c = &Container{Name: fmt.Sprintf("%v-%d-%d-", namePrefix, os.Getpid(), containerIdx)}

	// Start the container.
	log.WithField("container", c).Info("About to run container")
	runArgs := append([]string{"run", "--name", c.Name}, args...)
	runCmd := exec.Command("docker", runArgs...)
	err := runCmd.Start()
	Expect(err).NotTo(HaveOccurred())
	c.WaitRunning(10 * time.Second)

	// Remember that this container is now running.
	runningContainers = append(runningContainers, c)

	// Fill in rest of container struct.
	c.IP = c.GetIP()
	c.Hostname = c.GetHostname()
	c.Stop = func() {
		if !c.stopped {
			// We haven't previously tried to stop this container.
			c.stopped = true
			runCmd.Process.Signal(os.Interrupt)
			c.WaitNotRunning(10 * time.Second)

			// And now to be really sure that the container is cleaned up.
			utils.RunMayFail("docker", "rm", "-f", c.Name)
		}
	}
	c.binaries = set.New()
	log.WithField("container", c).Info("Container now running")
	return
}

func (c *Container) DockerInspect(format string) string {
	inspectCmd := exec.Command("docker", "inspect",
		"--format="+format,
		c.Name,
	)
	outputBytes, err := inspectCmd.CombinedOutput()
	Expect(err).NotTo(HaveOccurred())
	return string(outputBytes)
}

func (c *Container) GetIP() string {
	output := c.DockerInspect("{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}")
	return strings.TrimSpace(output)
}

func (c *Container) GetHostname() string {
	output := c.DockerInspect("{{.Config.Hostname}}")
	return strings.TrimSpace(output)
}

func (c *Container) WaitRunning(timeout time.Duration) {
	log.Info("Wait for container to be listed in docker ps")
	start := time.Now()
	for {
		cmd := exec.Command("docker", "ps")
		out, err := cmd.CombinedOutput()
		Expect(err).NotTo(HaveOccurred())
		if strings.Contains(string(out), c.Name) {
			break
		}
		if time.Since(start) > timeout {
			log.Panic("Timed out waiting for container to be listed.")
		}
	}
}

func (c *Container) WaitNotRunning(timeout time.Duration) {
	log.Info("Wait for container not to be listed in docker ps")
	start := time.Now()
	for {
		cmd := exec.Command("docker", "ps")
		out, err := cmd.CombinedOutput()
		Expect(err).NotTo(HaveOccurred())
		if !strings.Contains(string(out), c.Name) {
			break
		}
		if time.Since(start) > timeout {
			log.Panic("Timed out waiting for container not to be listed.")
		}
	}
}

var _ = AfterEach(func() {
	for _, c := range runningContainers {
		c.Stop()
	}
	runningContainers = []*Container{}
})

func (c *Container) EnsureBinary(name string) {
	if !c.binaries.Contains(name) {
		exec.Command("docker", "cp", "../bin/"+name, c.Name+":/"+name).Run()
		c.binaries.Add(name)
	}
}
<commit_msg>Allow 20 minutes for possible container image download<commit_after>// Copyright (c) 2017 Tigera, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package containers

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/projectcalico/felix/fv/utils"
	"github.com/projectcalico/libcalico-go/lib/set"
	log "github.com/sirupsen/logrus"
)

type Container struct {
	Name     string
	IP       string
	Hostname string
	Stop     func()
	stopped  bool
	binaries set.Set
}

var containerIdx = 0

var runningContainers = []*Container{}

func Run(namePrefix string, args ...string) (c *Container) {

	// Build unique container name and struct.
	containerIdx++
	c = &Container{Name: fmt.Sprintf("%v-%d-%d-", namePrefix, os.Getpid(), containerIdx)}

	// Start the container.
	log.WithField("container", c).Info("About to run container")
	runArgs := append([]string{"run", "--name", c.Name}, args...)
	runCmd := exec.Command("docker", runArgs...)
	err := runCmd.Start()
	Expect(err).NotTo(HaveOccurred())

	// It might take a very long time for the container to show as running, if the image needs
	// to be downloaded - e.g. when running on semaphore.
	c.WaitRunning(20 * 60 * time.Second)

	// Remember that this container is now running.
	runningContainers = append(runningContainers, c)

	// Fill in rest of container struct.
	c.IP = c.GetIP()
	c.Hostname = c.GetHostname()
	c.Stop = func() {
		if !c.stopped {
			// We haven't previously tried to stop this container.
			c.stopped = true
			runCmd.Process.Signal(os.Interrupt)
			c.WaitNotRunning(10 * time.Second)

			// And now to be really sure that the container is cleaned up.
			utils.RunMayFail("docker", "rm", "-f", c.Name)
		}
	}
	c.binaries = set.New()
	log.WithField("container", c).Info("Container now running")
	return
}

func (c *Container) DockerInspect(format string) string {
	inspectCmd := exec.Command("docker", "inspect",
		"--format="+format,
		c.Name,
	)
	outputBytes, err := inspectCmd.CombinedOutput()
	Expect(err).NotTo(HaveOccurred())
	return string(outputBytes)
}

func (c *Container) GetIP() string {
	output := c.DockerInspect("{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}")
	return strings.TrimSpace(output)
}

func (c *Container) GetHostname() string {
	output := c.DockerInspect("{{.Config.Hostname}}")
	return strings.TrimSpace(output)
}

func (c *Container) WaitRunning(timeout time.Duration) {
	log.Info("Wait for container to be listed in docker ps")
	start := time.Now()
	for {
		cmd := exec.Command("docker", "ps")
		out, err := cmd.CombinedOutput()
		Expect(err).NotTo(HaveOccurred())
		if strings.Contains(string(out), c.Name) {
			break
		}
		if time.Since(start) > timeout {
			log.Panic("Timed out waiting for container to be listed.")
		}
	}
}

func (c *Container) WaitNotRunning(timeout time.Duration) {
	log.Info("Wait for container not to be listed in docker ps")
	start := time.Now()
	for {
		cmd := exec.Command("docker", "ps")
		out, err := cmd.CombinedOutput()
		Expect(err).NotTo(HaveOccurred())
		if !strings.Contains(string(out), c.Name) {
			break
		}
		if time.Since(start) > timeout {
			log.Panic("Timed out waiting for container not to be listed.")
		}
	}
}

var _ = AfterEach(func() {
	for _, c := range runningContainers {
		c.Stop()
	}
	runningContainers = []*Container{}
})

func (c *Container) EnsureBinary(name string) {
	if !c.binaries.Contains(name) {
		exec.Command("docker", "cp", "../bin/"+name, c.Name+":/"+name).Run()
		c.binaries.Add(name)
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 gf Author(https://gitee.com/johng/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/johng/gf.

// Kafka Client.
package gkafka

import (
    "time"
    "strings"
    "github.com/Shopify/sarama"
    "github.com/bsm/sarama-cluster"
    "errors"
)

var (
    // 当使用Topics方法获取所有topic后，进行过滤忽略的topic，多个以','号分隔
    ignoreTopics = map[string]bool {
        "__consumer_offsets" : true,
    }
)

// kafka Client based on sarama.Config
type Config struct {
    GroupId        string // group id for consumer.
    Servers        string // server list, multiple servers joined by ','.
    Topics         string // topic list, multiple topics joined by ','.
    AutoMarkOffset bool   // auto mark message read after consumer message from server
    sarama.Config
}

// Kafka Client(Consumer/SyncProducer/AsyncProducer)
type Client struct {
    Config        *Config
    consumer      *cluster.Consumer
    rawConsumer   sarama.Consumer
    syncProducer  sarama.SyncProducer
    asyncProducer sarama.AsyncProducer
}

// Kafka Message.
type Message struct {
    Value          []byte
    Key            []byte
    Topic          string
    Partition      int
    Offset         int
    client         *Client
    consumerMsg    *sarama.ConsumerMessage
}

// New a kafka client.
func NewClient(config *Config) *Client {
    return &Client {
        Config : config,
    }
}

// New a default configuration object.
func NewConfig() *Config {
    config       := &Config{}
    config.Config = *sarama.NewConfig()

    // default config for consumer
    config.Consumer.Return.Errors          = true
    config.Consumer.Offsets.Initial        = sarama.OffsetOldest
    config.Consumer.Offsets.CommitInterval = 1 * time.Second

    // default config for producer
    config.Producer.Return.Errors          = true
    config.Producer.Return.Successes       = true
    config.Producer.Timeout                = 5 * time.Second

    config.AutoMarkOffset                  = true
    return config
}

// Close client.
func (client *Client) Close() {
    if client.consumer != nil {
        client.consumer.Close()
    }
    if client.syncProducer != nil {
        client.syncProducer.Close()
    }
    if client.asyncProducer != nil {
        client.asyncProducer.Close()
    }
}

// Get all topics from kafka server.
func (client *Client) Topics() ([]string, error) {
    if client.rawConsumer == nil {
        if c, err := sarama.NewConsumer(strings.Split(client.Config.Servers, ","), &client.Config.Config); err != nil {
            return nil, err
        } else {
            client.rawConsumer = c
        }
    }
    if topics, err := client.rawConsumer.Topics(); err == nil {
        for k, v := range topics {
            if _, ok := ignoreTopics[v]; ok {
                topics = append(topics[ : k], topics[k + 1 : ]...)
            }
        }
        return topics, nil
    } else {
        return nil, err
    }
}

// Receive message from kafka from specified topics in config, in BLOCKING way, gkafka will handle offset tracking automatically.
func (client *Client) Receive() (*Message, error) {
    if client.consumer == nil {
        config       := cluster.NewConfig()
        config.Config = client.Config.Config
        config.Group.Return.Notifications = false

        c, err := cluster.NewConsumer(strings.Split(client.Config.Servers, ","), client.Config.GroupId, strings.Split(client.Config.Topics, ","), config)
        if err != nil {
            return nil, err
        } else {
            client.consumer = c
        }
    }
    errorsChan  := client.consumer.Errors()
    notifyChan  := client.consumer.Notifications()
    messageChan := client.consumer.Messages()
    for {
        select {
            case msg := <- messageChan:
                if client.Config.AutoMarkOffset {
                    client.consumer.MarkOffset(msg, "")
                }
                return &Message {
                    Value       : msg.Value,
                    Key         : msg.Key,
                    Topic       : msg.Topic,
                    Partition   : int(msg.Partition),
                    Offset      : int(msg.Offset),
                    client      : client,
                    consumerMsg : msg,
                }, nil

            case err := <-errorsChan:
                if err != nil {
                    return nil, err
                }

            case <-notifyChan:
        }
    }

    return nil, errors.New("unknown error")
}

// Send data to kafka in synchronized way.
func (client *Client) SyncSend(message *Message) error {
    if client.syncProducer == nil {
        if p, err := sarama.NewSyncProducer(strings.Split(client.Config.Servers, ","), &client.Config.Config); err != nil {
            return err
        } else {
            client.syncProducer = p
        }
    }
    for _, topic := range strings.Split(client.Config.Topics, ",") {
        msg := messageToProducerMessage(message)
        msg.Topic = topic
        if _, _, err := client.syncProducer.SendMessage(msg); err != nil {
            return err
        }
    }
    return nil
}

// Send data to kafka in asynchronized way.
func (client *Client) AsyncSend(message *Message) error {
    if client.asyncProducer == nil {
        if p, err := sarama.NewAsyncProducer(strings.Split(client.Config.Servers, ","), &client.Config.Config); err != nil {
            return err
        } else {
            client.asyncProducer = p
            //go func(p sarama.AsyncProducer) {
            //    errors  := p.Errors()
            //    success := p.Successes()
            //    for {
            //        select {
            //            case err := <-errors:
            //                if err != nil {
            //                    glog.Error(err)
            //                }
            //            case <-success:
            //        }
            //    }
            //}(client.asyncProducer)
        }
    }

    for _, topic := range strings.Split(client.Config.Topics, ",") {
        msg      := messageToProducerMessage(message)
        msg.Topic = topic
        client.asyncProducer.Input() <- msg
    }
    return nil
}

// Convert *gkafka.Message to *sarama.ProducerMessage
func messageToProducerMessage(message *Message) *sarama.ProducerMessage {
    return &sarama.ProducerMessage {
        Topic     : message.Topic,
        Key       : sarama.ByteEncoder(message.Key),
        Value     : sarama.ByteEncoder(message.Value),
        Partition : int32(message.Partition),
        Offset    : int64(message.Offset),
    }
}
<commit_msg>修复gkafka对象关闭问题<commit_after>// Copyright 2018 gf Author(https://gitee.com/johng/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/johng/gf.

// Kafka Client.
package gkafka

import (
    "time"
    "strings"
    "github.com/Shopify/sarama"
    "github.com/bsm/sarama-cluster"
    "errors"
)

var (
    // 当使用Topics方法获取所有topic后，进行过滤忽略的topic，多个以','号分隔
    ignoreTopics = map[string]bool {
        "__consumer_offsets" : true,
    }
)

// kafka Client based on sarama.Config
type Config struct {
    GroupId        string // group id for consumer.
    Servers        string // server list, multiple servers joined by ','.
    Topics         string // topic list, multiple topics joined by ','.
    AutoMarkOffset bool   // auto mark message read after consumer message from server
    sarama.Config
}

// Kafka Client(Consumer/SyncProducer/AsyncProducer)
type Client struct {
    Config        *Config
    consumer      *cluster.Consumer
    rawConsumer   sarama.Consumer
    syncProducer  sarama.SyncProducer
    asyncProducer sarama.AsyncProducer
}

// Kafka Message.
type Message struct {
    Value          []byte
    Key            []byte
    Topic          string
    Partition      int
    Offset         int
    client         *Client
    consumerMsg    *sarama.ConsumerMessage
}

// New a kafka client.
func NewClient(config *Config) *Client {
    return &Client {
        Config : config,
    }
}

// New a default configuration object.
func NewConfig() *Config {
    config       := &Config{}
    config.Config = *sarama.NewConfig()

    // default config for consumer
    config.Consumer.Return.Errors          = true
    config.Consumer.Offsets.Initial        = sarama.OffsetOldest
    config.Consumer.Offsets.CommitInterval = 1 * time.Second

    // default config for producer
    config.Producer.Return.Errors          = true
    config.Producer.Return.Successes       = true
    config.Producer.Timeout                = 5 * time.Second

    config.AutoMarkOffset                  = true
    return config
}

// Close client.
func (client *Client) Close() {
    if client.rawConsumer != nil {
        client.rawConsumer.Close()
    }
    if client.consumer != nil {
        client.consumer.Close()
    }
    if client.syncProducer != nil {
        client.syncProducer.Close()
    }
    if client.asyncProducer != nil {
        client.asyncProducer.Close()
    }
}

// Get all topics from kafka server.
func (client *Client) Topics() ([]string, error) {
    if client.rawConsumer == nil {
        if c, err := sarama.NewConsumer(strings.Split(client.Config.Servers, ","), &client.Config.Config); err != nil {
            return nil, err
        } else {
            client.rawConsumer = c
        }
    }
    if topics, err := client.rawConsumer.Topics(); err == nil {
        for k, v := range topics {
            if _, ok := ignoreTopics[v]; ok {
                topics = append(topics[ : k], topics[k + 1 : ]...)
            }
        }
        return topics, nil
    } else {
        return nil, err
    }
}

// Receive message from kafka from specified topics in config, in BLOCKING way, gkafka will handle offset tracking automatically.
func (client *Client) Receive() (*Message, error) {
    if client.consumer == nil {
        config       := cluster.NewConfig()
        config.Config = client.Config.Config
        config.Group.Return.Notifications = false

        c, err := cluster.NewConsumer(strings.Split(client.Config.Servers, ","), client.Config.GroupId, strings.Split(client.Config.Topics, ","), config)
        if err != nil {
            return nil, err
        } else {
            client.consumer = c
        }
    }
    errorsChan  := client.consumer.Errors()
    notifyChan  := client.consumer.Notifications()
    messageChan := client.consumer.Messages()
    for {
        select {
            case msg := <- messageChan:
                if client.Config.AutoMarkOffset {
                    client.consumer.MarkOffset(msg, "")
                }
                return &Message {
                    Value       : msg.Value,
                    Key         : msg.Key,
                    Topic       : msg.Topic,
                    Partition   : int(msg.Partition),
                    Offset      : int(msg.Offset),
                    client      : client,
                    consumerMsg : msg,
                }, nil

            case err := <-errorsChan:
                if err != nil {
                    return nil, err
                }

            case <-notifyChan:
        }
    }

    return nil, errors.New("unknown error")
}

// Send data to kafka in synchronized way.
func (client *Client) SyncSend(message *Message) error {
    if client.syncProducer == nil {
        if p, err := sarama.NewSyncProducer(strings.Split(client.Config.Servers, ","), &client.Config.Config); err != nil {
            return err
        } else {
            client.syncProducer = p
        }
    }
    for _, topic := range strings.Split(client.Config.Topics, ",") {
        msg := messageToProducerMessage(message)
        msg.Topic = topic
        if _, _, err := client.syncProducer.SendMessage(msg); err != nil {
            return err
        }
    }
    return nil
}

// Send data to kafka in asynchronized way.
func (client *Client) AsyncSend(message *Message) error {
    if client.asyncProducer == nil {
        if p, err := sarama.NewAsyncProducer(strings.Split(client.Config.Servers, ","), &client.Config.Config); err != nil {
            return err
        } else {
            client.asyncProducer = p
            //go func(p sarama.AsyncProducer) {
            //    errors  := p.Errors()
            //    success := p.Successes()
            //    for {
            //        select {
            //            case err := <-errors:
            //                if err != nil {
            //                    glog.Error(err)
            //                }
            //            case <-success:
            //        }
            //    }
            //}(client.asyncProducer)
        }
    }

    for _, topic := range strings.Split(client.Config.Topics, ",") {
        msg      := messageToProducerMessage(message)
        msg.Topic = topic
        client.asyncProducer.Input() <- msg
    }
    return nil
}

// Convert *gkafka.Message to *sarama.ProducerMessage
func messageToProducerMessage(message *Message) *sarama.ProducerMessage {
    return &sarama.ProducerMessage {
        Topic     : message.Topic,
        Key       : sarama.ByteEncoder(message.Key),
        Value     : sarama.ByteEncoder(message.Value),
        Partition : int32(message.Partition),
        Offset    : int64(message.Offset),
    }
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"go/build"
	"os"
	"path/filepath"
	"sync"
)

type Package struct {
	Name  string
	Lines int
}

func plural(num int, str, p string) string {
	if num == 1 {
		return str
	}

	return str + p
}

var (
	wg   sync.WaitGroup
	done = NewDone()
)

func countLines(linesC chan<- Package, pkg *build.Package) {
	defer wg.Done()

	for _, ipath := range pkg.Imports {
		if done.Check(ipath) {
			continue
		}

		pkg, err := build.Import(ipath, ".", 0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to import %q: %v", ipath, err)
			os.Exit(1)
		}

		wg.Add(1)
		go countLines(linesC, pkg)
	}

	for _, file := range pkg.GoFiles {
		path := filepath.Join(pkg.Dir, file)

		wg.Add(1)
		go func(path string) {
			defer wg.Done()

			file, err := os.Open(path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Failed to open %q: %v", path, err)
				os.Exit(1)
			}
			defer file.Close()
		}(path)
	}
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %v [import path] ...\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()

	if flag.NArg() == 0 {
		flag.Usage()
		os.Exit(2)
	}

	linesC := make(chan Package)

	for _, ipath := range flag.Args() {
		if done.Check(ipath) {
			continue
		}

		pkg, err := build.Import(ipath, ".", 0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to import %q: %v", ipath, err)
			os.Exit(1)
		}

		wg.Add(1)
		go countLines(linesC, pkg)
	}

	go func() {
		wg.Wait()
		close(linesC)
	}()

	var total int
	for pkg := range linesC {
		fmt.Printf("%v: %v %v\n", pkg.Name, pkg.Lines, plural(pkg.Lines, "line", "s"))
		total += pkg.Lines
	}
	fmt.Printf("%v %v total.", total, plural(total, "line", "s"))
}
<commit_msg>It works.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"go/build"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
)

type Package struct {
	Name  string
	Lines int64
}

func plural(num int64, str, p string) string {
	if num == 1 {
		return str
	}

	return str + p
}

var (
	wg   sync.WaitGroup
	done = NewDone()
)

func countLines(linesC chan<- Package, pkg *build.Package) {
	defer wg.Done()

	for _, ipath := range pkg.Imports {
		if done.Check(ipath) {
			continue
		}

		pkg, err := build.Import(ipath, ".", 0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to import %q: %v", ipath, err)
			os.Exit(1)
		}

		wg.Add(1)
		go countLines(linesC, pkg)
	}

	var lines int64
	var wg sync.WaitGroup

	for _, file := range pkg.GoFiles {
		path := filepath.Join(pkg.Dir, file)

		wg.Add(1)
		go func(path string) {
			defer wg.Done()

			file, err := os.Open(path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Failed to open %q: %v", path, err)
				os.Exit(1)
			}
			defer file.Close()

			s := bufio.NewScanner(file)
			for s.Scan() {
				atomic.AddInt64(&lines, 1)
			}
		}(path)
	}

	wg.Wait()

	linesC <- Package{
		Name:  pkg.ImportPath,
		Lines: lines,
	}
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %v [import path] ...\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()

	if flag.NArg() == 0 {
		flag.Usage()
		os.Exit(2)
	}

	linesC := make(chan Package)

	for _, ipath := range flag.Args() {
		if done.Check(ipath) {
			continue
		}

		pkg, err := build.Import(ipath, ".", 0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to import %q: %v", ipath, err)
			os.Exit(1)
		}

		wg.Add(1)
		go countLines(linesC, pkg)
	}

	go func() {
		wg.Wait()
		close(linesC)
	}()

	var total int64
	for pkg := range linesC {
		fmt.Printf("%v: %v %v.\n", pkg.Name, pkg.Lines, plural(pkg.Lines, "line", "s"))
		total += pkg.Lines
	}
	fmt.Printf("%v %v total.\n", total, plural(total, "line", "s"))
}
<|endoftext|>


<commit_before>package lfs

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/github/git-lfs/git"
)

var (
	Config        = NewConfig()
	defaultRemote = "origin"
)

type Configuration struct {
	CurrentRemote         string
	httpClient            *HttpClient
	redirectingHttpClient *http.Client
	envVars               map[string]string
	isTracingHttp         bool
	isLoggingStats        bool

	loading   sync.Mutex // guards initialization of gitConfig and remotes
	gitConfig map[string]string
	remotes   []string
}

func NewConfig() *Configuration {
	c := &Configuration{
		CurrentRemote: defaultRemote,
		envVars:       make(map[string]string),
	}
	c.isTracingHttp = c.GetenvBool("GIT_CURL_VERBOSE", false)
	c.isLoggingStats = c.GetenvBool("GIT_LOG_STATS", false)
	return c
}

func (c *Configuration) Getenv(key string) string {
	if i, ok := c.envVars[key]; ok {
		return i
	}

	v := os.Getenv(key)
	c.envVars[key] = v
	return v
}

func (c *Configuration) Setenv(key, value string) error {
	// Check see if we have this in our cache, if so update it
	if _, ok := c.envVars[key]; ok {
		c.envVars[key] = value
	}

	// Now set in process
	return os.Setenv(key, value)
}

// GetenvBool parses a boolean environment variable and returns the result as a bool.
// If the environment variable is unset, empty, or if the parsing fails,
// the value of def (default) is returned instead.
func (c *Configuration) GetenvBool(key string, def bool) bool {
	s := c.Getenv(key)
	if len(s) == 0 {
		return def
	}

	b, err := strconv.ParseBool(s)
	if err != nil {
		return def
	}
	return b
}

func (c *Configuration) Endpoint() Endpoint {
	if url, ok := c.GitConfig("lfs.url"); ok {
		return NewEndpoint(url)
	}

	if len(c.CurrentRemote) > 0 && c.CurrentRemote != defaultRemote {
		if endpoint := c.RemoteEndpoint(c.CurrentRemote); len(endpoint.Url) > 0 {
			return endpoint
		}
	}

	return c.RemoteEndpoint(defaultRemote)
}

func (c *Configuration) ConcurrentTransfers() int {
	uploads := 3

	if v, ok := c.GitConfig("lfs.concurrenttransfers"); ok {
		n, err := strconv.Atoi(v)
		if err == nil && n > 0 {
			uploads = n
		}
	}

	return uploads
}

func (c *Configuration) BatchTransfer() bool {
	if v, ok := c.GitConfig("lfs.batch"); ok {
		if v == "true" || v == "" {
			return true
		}

		// Any numeric value except 0 is considered true
		if n, err := strconv.Atoi(v); err == nil && n != 0 {
			return true
		}
	}
	return false
}

// PrivateAccess will retrieve the access value and return true if
// the value is set to private. When a repo is marked as having private
// access, the http requests for the batch api will fetch the credentials
// before running, otherwise the request will run without credentials.
func (c *Configuration) PrivateAccess() bool {
	key := fmt.Sprintf("lfs.%s.access", c.Endpoint().Url)
	if v, ok := c.GitConfig(key); ok {
		if strings.ToLower(v) == "private" {
			return true
		}
	}
	return false
}

// SetPrivateAccess will set the private access flag in .git/config.
func (c *Configuration) SetPrivateAccess() {
	key := fmt.Sprintf("lfs.%s.access", c.Endpoint().Url)
	configFile := filepath.Join(LocalGitDir, "config")
	git.Config.SetLocal(configFile, key, "private")

	// Modify the config cache because it's checked again in this process
	// without being reloaded.
	c.loading.Lock()
	c.gitConfig[key] = "private"
	c.loading.Unlock()
}

func (c *Configuration) RemoteEndpoint(remote string) Endpoint {
	if len(remote) == 0 {
		remote = defaultRemote
	}
	if url, ok := c.GitConfig("remote." + remote + ".lfsurl"); ok {
		return NewEndpoint(url)
	}

	if url, ok := c.GitConfig("remote." + remote + ".url"); ok {
		return NewEndpointFromCloneURL(url)
	}

	return Endpoint{}
}

func (c *Configuration) Remotes() []string {
	c.loadGitConfig()
	return c.remotes
}

func (c *Configuration) GitConfig(key string) (string, bool) {
	c.loadGitConfig()
	value, ok := c.gitConfig[strings.ToLower(key)]
	return value, ok
}

func (c *Configuration) SetConfig(key, value string) {
	c.loadGitConfig()
	c.gitConfig[key] = value
}

func (c *Configuration) ObjectUrl(oid string) (*url.URL, error) {
	return ObjectUrl(c.Endpoint(), oid)
}

func (c *Configuration) loadGitConfig() {
	c.loading.Lock()
	defer c.loading.Unlock()

	if c.gitConfig != nil {
		return
	}

	uniqRemotes := make(map[string]bool)

	c.gitConfig = make(map[string]string)

	var output string
	listOutput, err := git.Config.List()
	if err != nil {
		panic(fmt.Errorf("Error listing git config: %s", err))
	}

	configFile := filepath.Join(LocalWorkingDir, ".gitconfig")
	fileOutput, err := git.Config.ListFromFile(configFile)
	if err != nil {
		panic(fmt.Errorf("Error listing git config from file: %s", err))
	}

	localConfig := filepath.Join(LocalGitDir, "config")
	localOutput, err := git.Config.ListFromFile(localConfig)
	if err != nil {
		panic(fmt.Errorf("Error listing git config from file %s", err))
	}

	output = fileOutput + "\n" + listOutput + "\n" + localOutput

	lines := strings.Split(output, "\n")
	for _, line := range lines {
		pieces := strings.SplitN(line, "=", 2)
		if len(pieces) < 2 {
			continue
		}
		key := strings.ToLower(pieces[0])
		c.gitConfig[key] = pieces[1]

		keyParts := strings.Split(key, ".")
		if len(keyParts) > 1 && keyParts[0] == "remote" {
			remote := keyParts[1]
			uniqRemotes[remote] = remote == "origin"
		}
	}

	c.remotes = make([]string, 0, len(uniqRemotes))
	for remote, isOrigin := range uniqRemotes {
		if isOrigin {
			continue
		}
		c.remotes = append(c.remotes, remote)
	}
}
<commit_msg>Add an extra tracer<commit_after>package lfs

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/github/git-lfs/git"
	"github.com/rubyist/tracerx"
)

var (
	Config        = NewConfig()
	defaultRemote = "origin"
)

type Configuration struct {
	CurrentRemote         string
	httpClient            *HttpClient
	redirectingHttpClient *http.Client
	envVars               map[string]string
	isTracingHttp         bool
	isLoggingStats        bool

	loading   sync.Mutex // guards initialization of gitConfig and remotes
	gitConfig map[string]string
	remotes   []string
}

func NewConfig() *Configuration {
	c := &Configuration{
		CurrentRemote: defaultRemote,
		envVars:       make(map[string]string),
	}
	c.isTracingHttp = c.GetenvBool("GIT_CURL_VERBOSE", false)
	c.isLoggingStats = c.GetenvBool("GIT_LOG_STATS", false)
	return c
}

func (c *Configuration) Getenv(key string) string {
	if i, ok := c.envVars[key]; ok {
		return i
	}

	v := os.Getenv(key)
	c.envVars[key] = v
	return v
}

func (c *Configuration) Setenv(key, value string) error {
	// Check see if we have this in our cache, if so update it
	if _, ok := c.envVars[key]; ok {
		c.envVars[key] = value
	}

	// Now set in process
	return os.Setenv(key, value)
}

// GetenvBool parses a boolean environment variable and returns the result as a bool.
// If the environment variable is unset, empty, or if the parsing fails,
// the value of def (default) is returned instead.
func (c *Configuration) GetenvBool(key string, def bool) bool {
	s := c.Getenv(key)
	if len(s) == 0 {
		return def
	}

	b, err := strconv.ParseBool(s)
	if err != nil {
		return def
	}
	return b
}

func (c *Configuration) Endpoint() Endpoint {
	if url, ok := c.GitConfig("lfs.url"); ok {
		return NewEndpoint(url)
	}

	if len(c.CurrentRemote) > 0 && c.CurrentRemote != defaultRemote {
		if endpoint := c.RemoteEndpoint(c.CurrentRemote); len(endpoint.Url) > 0 {
			return endpoint
		}
	}

	return c.RemoteEndpoint(defaultRemote)
}

func (c *Configuration) ConcurrentTransfers() int {
	uploads := 3

	if v, ok := c.GitConfig("lfs.concurrenttransfers"); ok {
		n, err := strconv.Atoi(v)
		if err == nil && n > 0 {
			uploads = n
		}
	}

	return uploads
}

func (c *Configuration) BatchTransfer() bool {
	if v, ok := c.GitConfig("lfs.batch"); ok {
		if v == "true" || v == "" {
			return true
		}

		// Any numeric value except 0 is considered true
		if n, err := strconv.Atoi(v); err == nil && n != 0 {
			return true
		}
	}
	return false
}

// PrivateAccess will retrieve the access value and return true if
// the value is set to private. When a repo is marked as having private
// access, the http requests for the batch api will fetch the credentials
// before running, otherwise the request will run without credentials.
func (c *Configuration) PrivateAccess() bool {
	key := fmt.Sprintf("lfs.%s.access", c.Endpoint().Url)
	if v, ok := c.GitConfig(key); ok {
		if strings.ToLower(v) == "private" {
			return true
		}
	}
	return false
}

// SetPrivateAccess will set the private access flag in .git/config.
func (c *Configuration) SetPrivateAccess() {
	tracerx.Printf("setting repository access to private")
	key := fmt.Sprintf("lfs.%s.access", c.Endpoint().Url)
	configFile := filepath.Join(LocalGitDir, "config")
	git.Config.SetLocal(configFile, key, "private")

	// Modify the config cache because it's checked again in this process
	// without being reloaded.
	c.loading.Lock()
	c.gitConfig[key] = "private"
	c.loading.Unlock()
}

func (c *Configuration) RemoteEndpoint(remote string) Endpoint {
	if len(remote) == 0 {
		remote = defaultRemote
	}
	if url, ok := c.GitConfig("remote." + remote + ".lfsurl"); ok {
		return NewEndpoint(url)
	}

	if url, ok := c.GitConfig("remote." + remote + ".url"); ok {
		return NewEndpointFromCloneURL(url)
	}

	return Endpoint{}
}

func (c *Configuration) Remotes() []string {
	c.loadGitConfig()
	return c.remotes
}

func (c *Configuration) GitConfig(key string) (string, bool) {
	c.loadGitConfig()
	value, ok := c.gitConfig[strings.ToLower(key)]
	return value, ok
}

func (c *Configuration) SetConfig(key, value string) {
	c.loadGitConfig()
	c.gitConfig[key] = value
}

func (c *Configuration) ObjectUrl(oid string) (*url.URL, error) {
	return ObjectUrl(c.Endpoint(), oid)
}

func (c *Configuration) loadGitConfig() {
	c.loading.Lock()
	defer c.loading.Unlock()

	if c.gitConfig != nil {
		return
	}

	uniqRemotes := make(map[string]bool)

	c.gitConfig = make(map[string]string)

	var output string
	listOutput, err := git.Config.List()
	if err != nil {
		panic(fmt.Errorf("Error listing git config: %s", err))
	}

	configFile := filepath.Join(LocalWorkingDir, ".gitconfig")
	fileOutput, err := git.Config.ListFromFile(configFile)
	if err != nil {
		panic(fmt.Errorf("Error listing git config from file: %s", err))
	}

	localConfig := filepath.Join(LocalGitDir, "config")
	localOutput, err := git.Config.ListFromFile(localConfig)
	if err != nil {
		panic(fmt.Errorf("Error listing git config from file %s", err))
	}

	output = fileOutput + "\n" + listOutput + "\n" + localOutput

	lines := strings.Split(output, "\n")
	for _, line := range lines {
		pieces := strings.SplitN(line, "=", 2)
		if len(pieces) < 2 {
			continue
		}
		key := strings.ToLower(pieces[0])
		c.gitConfig[key] = pieces[1]

		keyParts := strings.Split(key, ".")
		if len(keyParts) > 1 && keyParts[0] == "remote" {
			remote := keyParts[1]
			uniqRemotes[remote] = remote == "origin"
		}
	}

	c.remotes = make([]string, 0, len(uniqRemotes))
	for remote, isOrigin := range uniqRemotes {
		if isOrigin {
			continue
		}
		c.remotes = append(c.remotes, remote)
	}
}
<|endoftext|>


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

// Package leaktest provides tools to detect leaked goroutines in tests.
// To use it, call "defer leaktest.Check(t)()" at the beginning of each
// test that may use goroutines.
// copied out of the cockroachdb source tree with slight modifications to be
// more re-useable
package leaktest

import (
	"context"
	"fmt"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"time"
)

// TickerInterval defines the interval used by the ticker in Check* functions.
var TickerInterval = time.Millisecond * 50

type goroutine struct {
	id    uint64
	stack string
}

type goroutineByID []*goroutine

func (g goroutineByID) Len() int           { return len(g) }
func (g goroutineByID) Less(i, j int) bool { return g[i].id < g[j].id }
func (g goroutineByID) Swap(i, j int)      { g[i], g[j] = g[j], g[i] }

func interestingGoroutine(g string) (*goroutine, error) {
	sl := strings.SplitN(g, "\n", 2)
	if len(sl) != 2 {
		return nil, fmt.Errorf("error parsing stack: %q", g)
	}
	stack := strings.TrimSpace(sl[1])
	if strings.HasPrefix(stack, "testing.RunTests") {
		return nil, nil
	}

	if stack == "" ||
		// Ignore HTTP keep alives
		strings.Contains(stack, ").readLoop(") ||
		strings.Contains(stack, ").writeLoop(") ||
		// Below are the stacks ignored by the upstream leaktest code.
		strings.Contains(stack, "testing.Main(") ||
		strings.Contains(stack, "testing.(*T).Run(") ||
		strings.Contains(stack, "runtime.goexit") ||
		strings.Contains(stack, "created by runtime.gc") ||
		strings.Contains(stack, "interestingGoroutines") ||
		strings.Contains(stack, "runtime.MHeap_Scavenger") ||
		strings.Contains(stack, "signal.signal_recv") ||
		strings.Contains(stack, "sigterm.handler") ||
		strings.Contains(stack, "runtime_mcall") ||
		strings.Contains(stack, "goroutine in C code") {
		return nil, nil
	}

	// Parse the goroutine's ID from the header line.
	h := strings.SplitN(sl[0], " ", 3)
	if len(h) < 3 {
		return nil, fmt.Errorf("error parsing stack header: %q", sl[0])
	}
	id, err := strconv.ParseUint(h[1], 10, 64)
	if err != nil {
		return nil, fmt.Errorf("error parsing goroutine id: %s", err)
	}

	return &goroutine{id: id, stack: strings.TrimSpace(g)}, nil
}

// interestingGoroutines returns all goroutines we care about for the purpose
// of leak checking. It excludes testing or runtime ones.
func interestingGoroutines(t ErrorReporter) []*goroutine {
	buf := make([]byte, 2<<20)
	buf = buf[:runtime.Stack(buf, true)]
	var gs []*goroutine
	for _, g := range strings.Split(string(buf), "\n\n") {
		gr, err := interestingGoroutine(g)
		if err != nil {
			t.Errorf("leaktest: %s", err)
			continue
		} else if gr == nil {
			continue
		}
		gs = append(gs, gr)
	}
	sort.Sort(goroutineByID(gs))
	return gs
}

// ErrorReporter is a tiny subset of a testing.TB to make testing not such a
// massive pain
type ErrorReporter interface {
	Errorf(format string, args ...interface{})
}

// Check snapshots the currently-running goroutines and returns a
// function to be run at the end of tests to see whether any
// goroutines leaked, waiting up to 5 seconds in error conditions
func Check(t ErrorReporter) func() {
	return CheckTimeout(t, 5*time.Second)
}

// CheckTimeout is the same as Check, but with a configurable timeout
func CheckTimeout(t ErrorReporter, dur time.Duration) func() {
	ctx, cancel := context.WithCancel(context.Background())
	fn := CheckContext(ctx, t)
	return func() {
		timer := time.AfterFunc(dur, cancel)
		fn()
		// Remember to clean up the timer and context
		timer.Stop()
		cancel()
	}
}

// CheckContext is the same as Check, but uses a context.Context for
// cancellation and timeout control
func CheckContext(ctx context.Context, t ErrorReporter) func() {
	orig := map[uint64]bool{}
	for _, g := range interestingGoroutines(t) {
		orig[g.id] = true
	}
	return func() {
		var leaked []string
		// fast check if we have no leaks
		leaked = make([]string, 0)
		for _, g := range interestingGoroutines(t) {
			if !orig[g.id] {
				leaked = append(leaked, g.stack)
			}
		}
		if len(leaked) == 0 {
			return
		}
		ticker := time.NewTicker(TickerInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				leaked = make([]string, 0)
				for _, g := range interestingGoroutines(t) {
					if !orig[g.id] {
						leaked = append(leaked, g.stack)
					}
				}
				if len(leaked) == 0 {
					return
				}
				continue
			case <-ctx.Done():
				t.Errorf("leaktest: %v", ctx.Err())
			}
			break
		}

		for _, g := range leaked {
			t.Errorf("leaktest: leaked goroutine: %v", g)
		}
	}
}
<commit_msg>add leakedGoroutines function<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package leaktest provides tools to detect leaked goroutines in tests.
// To use it, call "defer leaktest.Check(t)()" at the beginning of each
// test that may use goroutines.
// copied out of the cockroachdb source tree with slight modifications to be
// more re-useable
package leaktest

import (
	"context"
	"fmt"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"time"
)

// TickerInterval defines the interval used by the ticker in Check* functions.
var TickerInterval = time.Millisecond * 50

type goroutine struct {
	id    uint64
	stack string
}

type goroutineByID []*goroutine

func (g goroutineByID) Len() int           { return len(g) }
func (g goroutineByID) Less(i, j int) bool { return g[i].id < g[j].id }
func (g goroutineByID) Swap(i, j int)      { g[i], g[j] = g[j], g[i] }

func interestingGoroutine(g string) (*goroutine, error) {
	sl := strings.SplitN(g, "\n", 2)
	if len(sl) != 2 {
		return nil, fmt.Errorf("error parsing stack: %q", g)
	}
	stack := strings.TrimSpace(sl[1])
	if strings.HasPrefix(stack, "testing.RunTests") {
		return nil, nil
	}

	if stack == "" ||
		// Ignore HTTP keep alives
		strings.Contains(stack, ").readLoop(") ||
		strings.Contains(stack, ").writeLoop(") ||
		// Below are the stacks ignored by the upstream leaktest code.
		strings.Contains(stack, "testing.Main(") ||
		strings.Contains(stack, "testing.(*T).Run(") ||
		strings.Contains(stack, "runtime.goexit") ||
		strings.Contains(stack, "created by runtime.gc") ||
		strings.Contains(stack, "interestingGoroutines") ||
		strings.Contains(stack, "runtime.MHeap_Scavenger") ||
		strings.Contains(stack, "signal.signal_recv") ||
		strings.Contains(stack, "sigterm.handler") ||
		strings.Contains(stack, "runtime_mcall") ||
		strings.Contains(stack, "goroutine in C code") {
		return nil, nil
	}

	// Parse the goroutine's ID from the header line.
	h := strings.SplitN(sl[0], " ", 3)
	if len(h) < 3 {
		return nil, fmt.Errorf("error parsing stack header: %q", sl[0])
	}
	id, err := strconv.ParseUint(h[1], 10, 64)
	if err != nil {
		return nil, fmt.Errorf("error parsing goroutine id: %s", err)
	}

	return &goroutine{id: id, stack: strings.TrimSpace(g)}, nil
}

// interestingGoroutines returns all goroutines we care about for the purpose
// of leak checking. It excludes testing or runtime ones.
func interestingGoroutines(t ErrorReporter) []*goroutine {
	buf := make([]byte, 2<<20)
	buf = buf[:runtime.Stack(buf, true)]
	var gs []*goroutine
	for _, g := range strings.Split(string(buf), "\n\n") {
		gr, err := interestingGoroutine(g)
		if err != nil {
			t.Errorf("leaktest: %s", err)
			continue
		} else if gr == nil {
			continue
		}
		gs = append(gs, gr)
	}
	sort.Sort(goroutineByID(gs))
	return gs
}

// leakedGoroutines returns all goroutines we are considering leaked and
// the boolean flag indicating if no leaks detected
func leakedGoroutines(orig map[uint64]bool, interesting []*goroutine) ([]string, bool) {
	leaked := make([]string, 0)
	flag := true
	for _, g := range interesting {
		if !orig[g.id] {
			leaked = append(leaked, g.stack)
			flag = false
		}
	}
	return leaked, flag
}

// ErrorReporter is a tiny subset of a testing.TB to make testing not such a
// massive pain
type ErrorReporter interface {
	Errorf(format string, args ...interface{})
}

// Check snapshots the currently-running goroutines and returns a
// function to be run at the end of tests to see whether any
// goroutines leaked, waiting up to 5 seconds in error conditions
func Check(t ErrorReporter) func() {
	return CheckTimeout(t, 5*time.Second)
}

// CheckTimeout is the same as Check, but with a configurable timeout
func CheckTimeout(t ErrorReporter, dur time.Duration) func() {
	ctx, cancel := context.WithCancel(context.Background())
	fn := CheckContext(ctx, t)
	return func() {
		timer := time.AfterFunc(dur, cancel)
		fn()
		// Remember to clean up the timer and context
		timer.Stop()
		cancel()
	}
}

// CheckContext is the same as Check, but uses a context.Context for
// cancellation and timeout control
func CheckContext(ctx context.Context, t ErrorReporter) func() {
	orig := map[uint64]bool{}
	for _, g := range interestingGoroutines(t) {
		orig[g.id] = true
	}
	return func() {
		var leaked []string
		var ok bool
		// fast check if we have no leaks
		if leaked, ok = leakedGoroutines(orig, interestingGoroutines(t)); ok {
			return
		}
		ticker := time.NewTicker(TickerInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				if leaked, ok = leakedGoroutines(orig, interestingGoroutines(t)); ok {
					return
				}
				continue
			case <-ctx.Done():
				t.Errorf("leaktest: %v", ctx.Err())
			}
			break
		}

		for _, g := range leaked {
			t.Errorf("leaktest: leaked goroutine: %v", g)
		}
	}
}
<|endoftext|>


<commit_before>package lfs

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
)

type CallbackReader struct {
	C         CopyCallback
	TotalSize int64
	ReadSize  int64
	io.Reader
}

type CopyCallback func(int64, int64, int) error

func (w *CallbackReader) Read(p []byte) (int, error) {
	n, err := w.Reader.Read(p)

	if n > 0 {
		w.ReadSize += int64(n)
	}

	if err == nil && w.C != nil {
		err = w.C(w.TotalSize, w.ReadSize, n)
	}

	return n, err
}

func CopyWithCallback(writer io.Writer, reader io.Reader, totalSize int64, cb CopyCallback) (int64, error) {
	if cb == nil {
		return io.Copy(writer, reader)
	}

	cbReader := &CallbackReader{
		C:         cb,
		TotalSize: totalSize,
		Reader:    reader,
	}
	return io.Copy(writer, cbReader)
}

func CopyCallbackFile(event, filename string, index, totalFiles int) (CopyCallback, *os.File, error) {
	logPath := Config.Getenv("GIT_LFS_PROGRESS")
	if len(logPath) == 0 || len(filename) == 0 || len(event) == 0 {
		return nil, nil, nil
	}

	if !filepath.IsAbs(logPath) {
		return nil, nil, fmt.Errorf("GIT_LFS_PROGRESS must be an absolute path")
	}

	cbDir := filepath.Dir(logPath)
	if err := os.MkdirAll(cbDir, 0755); err != nil {
		return nil, nil, wrapProgressError(err, event, logPath)
	}

	file, err := os.OpenFile(logPath, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		return nil, file, wrapProgressError(err, event, logPath)
	}

	var prevWritten int64

	cb := CopyCallback(func(total int64, written int64, current int) error {
		if written != prevWritten {
			_, err := file.Write([]byte(fmt.Sprintf("%s %d/%d %d/%d %s\n", event, index, totalFiles, written, total, filename)))
			file.Sync()
			prevWritten = written
			return wrapProgressError(err, event, logPath)
		}

		return nil
	})

	return cb, file, nil
}

func wrapProgressError(err error, event, filename string) error {
	if err != nil {
		return fmt.Errorf("Error writing Git LFS %s progress to %s: %s", event, filename, err.Error())
	}

	return nil
}
<commit_msg>ラララララ ラー ララララー<commit_after>package lfs

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
)

type CallbackReader struct {
	C         CopyCallback
	TotalSize int64
	ReadSize  int64
	io.Reader
}

type CopyCallback func(totalSize int64, readSoFar int64, readSinceLast int) error

func (w *CallbackReader) Read(p []byte) (int, error) {
	n, err := w.Reader.Read(p)

	if n > 0 {
		w.ReadSize += int64(n)
	}

	if err == nil && w.C != nil {
		err = w.C(w.TotalSize, w.ReadSize, n)
	}

	return n, err
}

func CopyWithCallback(writer io.Writer, reader io.Reader, totalSize int64, cb CopyCallback) (int64, error) {
	if cb == nil {
		return io.Copy(writer, reader)
	}

	cbReader := &CallbackReader{
		C:         cb,
		TotalSize: totalSize,
		Reader:    reader,
	}
	return io.Copy(writer, cbReader)
}

func CopyCallbackFile(event, filename string, index, totalFiles int) (CopyCallback, *os.File, error) {
	logPath := Config.Getenv("GIT_LFS_PROGRESS")
	if len(logPath) == 0 || len(filename) == 0 || len(event) == 0 {
		return nil, nil, nil
	}

	if !filepath.IsAbs(logPath) {
		return nil, nil, fmt.Errorf("GIT_LFS_PROGRESS must be an absolute path")
	}

	cbDir := filepath.Dir(logPath)
	if err := os.MkdirAll(cbDir, 0755); err != nil {
		return nil, nil, wrapProgressError(err, event, logPath)
	}

	file, err := os.OpenFile(logPath, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		return nil, file, wrapProgressError(err, event, logPath)
	}

	var prevWritten int64

	cb := CopyCallback(func(total int64, written int64, current int) error {
		if written != prevWritten {
			_, err := file.Write([]byte(fmt.Sprintf("%s %d/%d %d/%d %s\n", event, index, totalFiles, written, total, filename)))
			file.Sync()
			prevWritten = written
			return wrapProgressError(err, event, logPath)
		}

		return nil
	})

	return cb, file, nil
}

func wrapProgressError(err error, event, filename string) error {
	if err != nil {
		return fmt.Errorf("Error writing Git LFS %s progress to %s: %s", event, filename, err.Error())
	}

	return nil
}
<|endoftext|>


<commit_before>package gateway

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"time"

	"github.com/docker/distribution/reference"
	"github.com/moby/buildkit/cache"
	"github.com/moby/buildkit/client/llb"
	"github.com/moby/buildkit/executor"
	"github.com/moby/buildkit/frontend"
	pb "github.com/moby/buildkit/frontend/gateway/pb"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/session"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
)

const (
	keySource           = "source"
	keyDevel            = "gateway-devel"
	exporterImageConfig = "containerimage.config"
)

func NewGatewayFrontend() frontend.Frontend {
	return &gatewayFrontend{}
}

type gatewayFrontend struct {
}

func filterPrefix(opts map[string]string, pfx string) map[string]string {
	m := map[string]string{}
	for k, v := range opts {
		if strings.HasPrefix(k, pfx) {
			m[strings.TrimPrefix(k, pfx)] = v
		}
	}
	return m
}

func (gf *gatewayFrontend) Solve(ctx context.Context, llbBridge frontend.FrontendLLBBridge, opts map[string]string) (retRef cache.ImmutableRef, exporterAttr map[string][]byte, retErr error) {

	source, ok := opts[keySource]
	if !ok {
		return nil, nil, errors.Errorf("no source specified for gateway")
	}

	sid := session.FromContext(ctx)

	_, isDevel := opts[keyDevel]
	var img ocispec.Image
	var rootFS cache.ImmutableRef

	if isDevel {
		ref, exp, err := llbBridge.Solve(session.NewContext(ctx, "gateway:"+sid),
			frontend.SolveRequest{
				Frontend:    source,
				FrontendOpt: filterPrefix(opts, "gateway-"),
			})
		if err != nil {
			return nil, nil, err
		}
		defer ref.Release(context.TODO())
		rootFS = ref
		config, ok := exp[exporterImageConfig]
		if ok {
			if err := json.Unmarshal(config, &img); err != nil {
				return nil, nil, err
			}
		}
	} else {
		sourceRef, err := reference.ParseNormalizedNamed(source)
		if err != nil {
			return nil, nil, err
		}

		dgst, config, err := llbBridge.ResolveImageConfig(ctx, sourceRef.String())
		if err != nil {
			return nil, nil, err
		}

		if err := json.Unmarshal(config, &img); err != nil {
			return nil, nil, err
		}

		sourceRef, err = reference.WithDigest(sourceRef, dgst)
		if err != nil {
			return nil, nil, err
		}

		src := llb.Image(sourceRef.String())

		def, err := src.Marshal()
		if err != nil {
			return nil, nil, err
		}

		ref, _, err := llbBridge.Solve(ctx, frontend.SolveRequest{
			Definition: def.ToPB(),
		})
		if err != nil {
			return nil, nil, err
		}
		defer ref.Release(context.TODO())
		rootFS = ref
	}

	lbf, err := newLLBBrideForwarder(ctx, llbBridge)
	defer lbf.conn.Close()
	if err != nil {
		return nil, nil, err
	}

	args := []string{"/run"}
	env := []string{}
	cwd := "/"
	if img.Config.Env != nil {
		env = img.Config.Env
	}
	if img.Config.Entrypoint != nil {
		args = img.Config.Entrypoint
	}
	if img.Config.WorkingDir != "" {
		cwd = img.Config.WorkingDir
	}
	i := 0
	for k, v := range opts {
		env = append(env, fmt.Sprintf("BUILDKIT_FRONTEND_OPT_%d", i)+"="+k+"="+v)
		i++
	}

	env = append(env, "BUILDKIT_SESSION_ID="+sid)

	defer func() {
		for _, r := range lbf.refs {
			if r != nil && (lbf.lastRef != r || retErr != nil) {
				r.Release(context.TODO())
			}
		}
	}()

	err = llbBridge.Exec(ctx, executor.Meta{
		Env:  env,
		Args: args,
		Cwd:  cwd,
	}, rootFS, lbf.Stdin, lbf.Stdout, os.Stderr)

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

	return lbf.lastRef, lbf.exporterAttr, nil
}

func newLLBBrideForwarder(ctx context.Context, llbBridge frontend.FrontendLLBBridge) (*llbBrideForwarder, error) {
	lbf := &llbBrideForwarder{
		llbBridge: llbBridge,
		refs:      map[string]cache.ImmutableRef{},
		pipe:      newPipe(),
	}

	server := grpc.NewServer()
	grpc_health_v1.RegisterHealthServer(server, health.NewServer())
	pb.RegisterLLBBridgeServer(server, lbf)

	go serve(ctx, server, lbf.conn)

	return lbf, nil
}

type pipe struct {
	Stdin  io.ReadCloser
	Stdout io.WriteCloser
	conn   net.Conn
}

func newPipe() *pipe {
	pr1, pw1, _ := os.Pipe()
	pr2, pw2, _ := os.Pipe()
	return &pipe{
		Stdin:  pr1,
		Stdout: pw2,
		conn: &conn{
			Reader: pr2,
			Writer: pw1,
			Closer: pw2,
		},
	}
}

type conn struct {
	io.Reader
	io.Writer
	io.Closer
}

func (s *conn) LocalAddr() net.Addr {
	return dummyAddr{}
}
func (s *conn) RemoteAddr() net.Addr {
	return dummyAddr{}
}
func (s *conn) SetDeadline(t time.Time) error {
	return nil
}
func (s *conn) SetReadDeadline(t time.Time) error {
	return nil
}
func (s *conn) SetWriteDeadline(t time.Time) error {
	return nil
}

type dummyAddr struct {
}

func (d dummyAddr) Network() string {
	return "pipe"
}

func (d dummyAddr) String() string {
	return "localhost"
}

type llbBrideForwarder struct {
	llbBridge    frontend.FrontendLLBBridge
	refs         map[string]cache.ImmutableRef
	lastRef      cache.ImmutableRef
	exporterAttr map[string][]byte
	*pipe
}

func (lbf *llbBrideForwarder) ResolveImageConfig(ctx context.Context, req *pb.ResolveImageConfigRequest) (*pb.ResolveImageConfigResponse, error) {
	dgst, dt, err := lbf.llbBridge.ResolveImageConfig(ctx, req.Ref)
	if err != nil {
		return nil, err
	}
	return &pb.ResolveImageConfigResponse{
		Digest: dgst,
		Config: dt,
	}, nil
}

func (lbf *llbBrideForwarder) Solve(ctx context.Context, req *pb.SolveRequest) (*pb.SolveResponse, error) {
	ref, expResp, err := lbf.llbBridge.Solve(ctx, frontend.SolveRequest{
		Definition: req.Definition,
		Frontend:   req.Frontend,
	})
	if err != nil {
		return nil, err
	}

	exp := map[string][]byte{}
	if err := json.Unmarshal(req.ExporterAttr, &exp); err != nil {
		return nil, err
	}

	if expResp != nil {
		for k, v := range expResp {
			exp[k] = v
		}
	}

	id := identity.NewID()
	lbf.refs[id] = ref
	if req.Final {
		lbf.lastRef = ref
		lbf.exporterAttr = exp
	}
	if ref == nil {
		id = ""
	}
	return &pb.SolveResponse{Ref: id}, nil
}
func (lbf *llbBrideForwarder) ReadFile(ctx context.Context, req *pb.ReadFileRequest) (*pb.ReadFileResponse, error) {
	ref, ok := lbf.refs[req.Ref]
	if !ok {
		return nil, errors.Errorf("no such ref: %v", req.Ref)
	}
	if ref == nil {
		return nil, errors.Wrapf(os.ErrNotExist, "%s no found", req.FilePath)
	}
	dt, err := cache.ReadFile(ctx, ref, req.FilePath)
	if err != nil {
		return nil, err
	}

	return &pb.ReadFileResponse{Data: dt}, nil
}

func (lbf *llbBrideForwarder) Ping(context.Context, *pb.PingRequest) (*pb.PongResponse, error) {
	return &pb.PongResponse{}, nil
}

func serve(ctx context.Context, grpcServer *grpc.Server, conn net.Conn) {
	go func() {
		<-ctx.Done()
		conn.Close()
	}()
	logrus.Debugf("serving grpc connection")
	(&http2.Server{}).ServeConn(conn, &http2.ServeConnOpts{Handler: grpcServer})
}
<commit_msg>gateway: allow skipping tag in source parameter<commit_after>package gateway

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"time"

	"github.com/docker/distribution/reference"
	"github.com/moby/buildkit/cache"
	"github.com/moby/buildkit/client/llb"
	"github.com/moby/buildkit/executor"
	"github.com/moby/buildkit/frontend"
	pb "github.com/moby/buildkit/frontend/gateway/pb"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/session"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
)

const (
	keySource           = "source"
	keyDevel            = "gateway-devel"
	exporterImageConfig = "containerimage.config"
)

func NewGatewayFrontend() frontend.Frontend {
	return &gatewayFrontend{}
}

type gatewayFrontend struct {
}

func filterPrefix(opts map[string]string, pfx string) map[string]string {
	m := map[string]string{}
	for k, v := range opts {
		if strings.HasPrefix(k, pfx) {
			m[strings.TrimPrefix(k, pfx)] = v
		}
	}
	return m
}

func (gf *gatewayFrontend) Solve(ctx context.Context, llbBridge frontend.FrontendLLBBridge, opts map[string]string) (retRef cache.ImmutableRef, exporterAttr map[string][]byte, retErr error) {

	source, ok := opts[keySource]
	if !ok {
		return nil, nil, errors.Errorf("no source specified for gateway")
	}

	sid := session.FromContext(ctx)

	_, isDevel := opts[keyDevel]
	var img ocispec.Image
	var rootFS cache.ImmutableRef

	if isDevel {
		ref, exp, err := llbBridge.Solve(session.NewContext(ctx, "gateway:"+sid),
			frontend.SolveRequest{
				Frontend:    source,
				FrontendOpt: filterPrefix(opts, "gateway-"),
			})
		if err != nil {
			return nil, nil, err
		}
		defer ref.Release(context.TODO())
		rootFS = ref
		config, ok := exp[exporterImageConfig]
		if ok {
			if err := json.Unmarshal(config, &img); err != nil {
				return nil, nil, err
			}
		}
	} else {
		sourceRef, err := reference.ParseNormalizedNamed(source)
		if err != nil {
			return nil, nil, err
		}

		dgst, config, err := llbBridge.ResolveImageConfig(ctx, reference.TagNameOnly(sourceRef).String())
		if err != nil {
			return nil, nil, err
		}

		if err := json.Unmarshal(config, &img); err != nil {
			return nil, nil, err
		}

		sourceRef, err = reference.WithDigest(sourceRef, dgst)
		if err != nil {
			return nil, nil, err
		}

		src := llb.Image(sourceRef.String())

		def, err := src.Marshal()
		if err != nil {
			return nil, nil, err
		}

		ref, _, err := llbBridge.Solve(ctx, frontend.SolveRequest{
			Definition: def.ToPB(),
		})
		if err != nil {
			return nil, nil, err
		}
		defer ref.Release(context.TODO())
		rootFS = ref
	}

	lbf, err := newLLBBrideForwarder(ctx, llbBridge)
	defer lbf.conn.Close()
	if err != nil {
		return nil, nil, err
	}

	args := []string{"/run"}
	env := []string{}
	cwd := "/"
	if img.Config.Env != nil {
		env = img.Config.Env
	}
	if img.Config.Entrypoint != nil {
		args = img.Config.Entrypoint
	}
	if img.Config.WorkingDir != "" {
		cwd = img.Config.WorkingDir
	}
	i := 0
	for k, v := range opts {
		env = append(env, fmt.Sprintf("BUILDKIT_FRONTEND_OPT_%d", i)+"="+k+"="+v)
		i++
	}

	env = append(env, "BUILDKIT_SESSION_ID="+sid)

	defer func() {
		for _, r := range lbf.refs {
			if r != nil && (lbf.lastRef != r || retErr != nil) {
				r.Release(context.TODO())
			}
		}
	}()

	err = llbBridge.Exec(ctx, executor.Meta{
		Env:  env,
		Args: args,
		Cwd:  cwd,
	}, rootFS, lbf.Stdin, lbf.Stdout, os.Stderr)

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

	return lbf.lastRef, lbf.exporterAttr, nil
}

func newLLBBrideForwarder(ctx context.Context, llbBridge frontend.FrontendLLBBridge) (*llbBrideForwarder, error) {
	lbf := &llbBrideForwarder{
		llbBridge: llbBridge,
		refs:      map[string]cache.ImmutableRef{},
		pipe:      newPipe(),
	}

	server := grpc.NewServer()
	grpc_health_v1.RegisterHealthServer(server, health.NewServer())
	pb.RegisterLLBBridgeServer(server, lbf)

	go serve(ctx, server, lbf.conn)

	return lbf, nil
}

type pipe struct {
	Stdin  io.ReadCloser
	Stdout io.WriteCloser
	conn   net.Conn
}

func newPipe() *pipe {
	pr1, pw1, _ := os.Pipe()
	pr2, pw2, _ := os.Pipe()
	return &pipe{
		Stdin:  pr1,
		Stdout: pw2,
		conn: &conn{
			Reader: pr2,
			Writer: pw1,
			Closer: pw2,
		},
	}
}

type conn struct {
	io.Reader
	io.Writer
	io.Closer
}

func (s *conn) LocalAddr() net.Addr {
	return dummyAddr{}
}
func (s *conn) RemoteAddr() net.Addr {
	return dummyAddr{}
}
func (s *conn) SetDeadline(t time.Time) error {
	return nil
}
func (s *conn) SetReadDeadline(t time.Time) error {
	return nil
}
func (s *conn) SetWriteDeadline(t time.Time) error {
	return nil
}

type dummyAddr struct {
}

func (d dummyAddr) Network() string {
	return "pipe"
}

func (d dummyAddr) String() string {
	return "localhost"
}

type llbBrideForwarder struct {
	llbBridge    frontend.FrontendLLBBridge
	refs         map[string]cache.ImmutableRef
	lastRef      cache.ImmutableRef
	exporterAttr map[string][]byte
	*pipe
}

func (lbf *llbBrideForwarder) ResolveImageConfig(ctx context.Context, req *pb.ResolveImageConfigRequest) (*pb.ResolveImageConfigResponse, error) {
	dgst, dt, err := lbf.llbBridge.ResolveImageConfig(ctx, req.Ref)
	if err != nil {
		return nil, err
	}
	return &pb.ResolveImageConfigResponse{
		Digest: dgst,
		Config: dt,
	}, nil
}

func (lbf *llbBrideForwarder) Solve(ctx context.Context, req *pb.SolveRequest) (*pb.SolveResponse, error) {
	ref, expResp, err := lbf.llbBridge.Solve(ctx, frontend.SolveRequest{
		Definition: req.Definition,
		Frontend:   req.Frontend,
	})
	if err != nil {
		return nil, err
	}

	exp := map[string][]byte{}
	if err := json.Unmarshal(req.ExporterAttr, &exp); err != nil {
		return nil, err
	}

	if expResp != nil {
		for k, v := range expResp {
			exp[k] = v
		}
	}

	id := identity.NewID()
	lbf.refs[id] = ref
	if req.Final {
		lbf.lastRef = ref
		lbf.exporterAttr = exp
	}
	if ref == nil {
		id = ""
	}
	return &pb.SolveResponse{Ref: id}, nil
}
func (lbf *llbBrideForwarder) ReadFile(ctx context.Context, req *pb.ReadFileRequest) (*pb.ReadFileResponse, error) {
	ref, ok := lbf.refs[req.Ref]
	if !ok {
		return nil, errors.Errorf("no such ref: %v", req.Ref)
	}
	if ref == nil {
		return nil, errors.Wrapf(os.ErrNotExist, "%s no found", req.FilePath)
	}
	dt, err := cache.ReadFile(ctx, ref, req.FilePath)
	if err != nil {
		return nil, err
	}

	return &pb.ReadFileResponse{Data: dt}, nil
}

func (lbf *llbBrideForwarder) Ping(context.Context, *pb.PingRequest) (*pb.PongResponse, error) {
	return &pb.PongResponse{}, nil
}

func serve(ctx context.Context, grpcServer *grpc.Server, conn net.Conn) {
	go func() {
		<-ctx.Done()
		conn.Close()
	}()
	logrus.Debugf("serving grpc connection")
	(&http2.Server{}).ServeConn(conn, &http2.ServeConnOpts{Handler: grpcServer})
}
<|endoftext|>