<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 pe implements access to PE (Microsoft Windows Portable Executable) files.
package pe

import (
	"debug/dwarf"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"strconv"
)

// A File represents an open PE file.
type File struct {
	FileHeader
	Sections []*Section

	closer io.Closer
}

type SectionHeader struct {
	Name                 string
	VirtualSize          uint32
	VirtualAddress       uint32
	Size                 uint32
	Offset               uint32
	PointerToRelocations uint32
	PointerToLineNumbers uint32
	NumberOfRelocations  uint16
	NumberOfLineNumbers  uint16
	Characteristics      uint32
}

type Section struct {
	SectionHeader

	// Embed ReaderAt for ReadAt method.
	// Do not embed SectionReader directly
	// to avoid having Read and Seek.
	// If a client wants Read and Seek it must use
	// Open() to avoid fighting over the seek offset
	// with other clients.
	io.ReaderAt
	sr *io.SectionReader
}

type ImportDirectory struct {
	OriginalFirstThunk uint32
	TimeDateStamp      uint32
	ForwarderChain     uint32
	Name               uint32
	FirstThunk         uint32

	dll string
}

// Data reads and returns the contents of the PE section.
func (s *Section) Data() ([]byte, os.Error) {
	dat := make([]byte, s.sr.Size())
	n, err := s.sr.ReadAt(dat, 0)
	return dat[0:n], err
}

// Open returns a new ReadSeeker reading the PE section.
func (s *Section) Open() io.ReadSeeker { return io.NewSectionReader(s.sr, 0, 1<<63-1) }

type FormatError struct {
	off int64
	msg string
	val interface{}
}

func (e *FormatError) String() string {
	msg := e.msg
	if e.val != nil {
		msg += fmt.Sprintf(" '%v'", e.val)
	}
	msg += fmt.Sprintf(" in record at byte %#x", e.off)
	return msg
}

// Open opens the named file using os.Open and prepares it for use as a PE binary.
func Open(name string) (*File, os.Error) {
	f, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	ff, err := NewFile(f)
	if err != nil {
		f.Close()
		return nil, err
	}
	ff.closer = f
	return ff, nil
}

// Close closes the File.
// If the File was created using NewFile directly instead of Open,
// Close has no effect.
func (f *File) Close() os.Error {
	var err os.Error
	if f.closer != nil {
		err = f.closer.Close()
		f.closer = nil
	}
	return err
}

// NewFile creates a new File for accessing a PE binary in an underlying reader.
func NewFile(r io.ReaderAt) (*File, os.Error) {
	f := new(File)
	sr := io.NewSectionReader(r, 0, 1<<63-1)

	var dosheader [96]byte
	if _, err := r.ReadAt(dosheader[0:], 0); err != nil {
		return nil, err
	}
	var base int64
	if dosheader[0] == 'M' && dosheader[1] == 'Z' {
		var sign [4]byte
		r.ReadAt(sign[0:], int64(dosheader[0x3c]))
		if !(sign[0] == 'P' && sign[1] == 'E' && sign[2] == 0 && sign[3] == 0) {
			return nil, os.NewError("Invalid PE File Format.")
		}
		base = int64(dosheader[0x3c]) + 4
	} else {
		base = int64(0)
	}
	sr.Seek(base, os.SEEK_SET)
	if err := binary.Read(sr, binary.LittleEndian, &f.FileHeader); err != nil {
		return nil, err
	}
	if f.FileHeader.Machine != IMAGE_FILE_MACHINE_UNKNOWN && f.FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64 && f.FileHeader.Machine != IMAGE_FILE_MACHINE_I386 {
		return nil, os.NewError("Invalid PE File Format.")
	}
	// get symbol string table
	sr.Seek(int64(f.FileHeader.PointerToSymbolTable+18*f.FileHeader.NumberOfSymbols), os.SEEK_SET)
	var l uint32
	if err := binary.Read(sr, binary.LittleEndian, &l); err != nil {
		return nil, err
	}
	ss := make([]byte, l)
	if _, err := r.ReadAt(ss, int64(f.FileHeader.PointerToSymbolTable+18*f.FileHeader.NumberOfSymbols)); err != nil {
		return nil, err
	}
	sr.Seek(base, os.SEEK_SET)
	binary.Read(sr, binary.LittleEndian, &f.FileHeader)
	sr.Seek(int64(f.FileHeader.SizeOfOptionalHeader), os.SEEK_CUR) //Skip OptionalHeader
	f.Sections = make([]*Section, f.FileHeader.NumberOfSections)
	for i := 0; i < int(f.FileHeader.NumberOfSections); i++ {
		sh := new(SectionHeader32)
		if err := binary.Read(sr, binary.LittleEndian, sh); err != nil {
			return nil, err
		}
		var name string
		if sh.Name[0] == '\x2F' {
			si, _ := strconv.Atoi(cstring(sh.Name[1:]))
			name, _ = getString(ss, si)
		} else {
			name = cstring(sh.Name[0:])
		}
		s := new(Section)
		s.SectionHeader = SectionHeader{
			Name:                 name,
			VirtualSize:          uint32(sh.VirtualSize),
			VirtualAddress:       uint32(sh.VirtualAddress),
			Size:                 uint32(sh.SizeOfRawData),
			Offset:               uint32(sh.PointerToRawData),
			PointerToRelocations: uint32(sh.PointerToRelocations),
			PointerToLineNumbers: uint32(sh.PointerToLineNumbers),
			NumberOfRelocations:  uint16(sh.NumberOfRelocations),
			NumberOfLineNumbers:  uint16(sh.NumberOfLineNumbers),
			Characteristics:      uint32(sh.Characteristics),
		}
		s.sr = io.NewSectionReader(r, int64(s.SectionHeader.Offset), int64(s.SectionHeader.Size))
		s.ReaderAt = s.sr
		f.Sections[i] = s
	}
	return f, nil
}

func cstring(b []byte) string {
	var i int
	for i = 0; i < len(b) && b[i] != 0; i++ {
	}
	return string(b[0:i])
}

// getString extracts a string from symbol string table.
func getString(section []byte, start int) (string, bool) {
	if start < 0 || start >= len(section) {
		return "", false
	}

	for end := start; end < len(section); end++ {
		if section[end] == 0 {
			return string(section[start:end]), true
		}
	}
	return "", false
}

// Section returns the first section with the given name, or nil if no such
// section exists.
func (f *File) Section(name string) *Section {
	for _, s := range f.Sections {
		if s.Name == name {
			return s
		}
	}
	return nil
}

func (f *File) DWARF() (*dwarf.Data, os.Error) {
	// There are many other DWARF sections, but these
	// are the required ones, and the debug/dwarf package
	// does not use the others, so don't bother loading them.
	var names = [...]string{"abbrev", "info", "str"}
	var dat [len(names)][]byte
	for i, name := range names {
		name = ".debug_" + name
		s := f.Section(name)
		if s == nil {
			continue
		}
		b, err := s.Data()
		if err != nil && uint32(len(b)) < s.Size {
			return nil, err
		}
		dat[i] = b
	}

	abbrev, info, str := dat[0], dat[1], dat[2]
	return dwarf.New(abbrev, nil, nil, info, nil, nil, nil, str)
}

// ImportedSymbols returns the names of all symbols
// referred to by the binary f that are expected to be
// satisfied by other libraries at dynamic load time.
// It does not return weak symbols.
func (f *File) ImportedSymbols() ([]string, os.Error) {
	ds := f.Section(".idata")
	if ds == nil {
		// not dynamic, so no libraries
		return nil, nil
	}
	d, err := ds.Data()
	if err != nil {
		return nil, err
	}
	var ida []ImportDirectory
	for len(d) > 0 {
		var dt ImportDirectory
		dt.OriginalFirstThunk = binary.LittleEndian.Uint32(d[0:4])
		dt.Name = binary.LittleEndian.Uint32(d[12:16])
		dt.FirstThunk = binary.LittleEndian.Uint32(d[16:20])
		d = d[20:]
		if dt.OriginalFirstThunk == 0 {
			break
		}
		ida = append(ida, dt)
	}
	names, _ := ds.Data()
	var all []string
	for _, dt := range ida {
		dt.dll, _ = getString(names, int(dt.Name-ds.VirtualAddress))
		d, _ = ds.Data()
		// seek to OriginalFirstThunk
		d = d[dt.OriginalFirstThunk-ds.VirtualAddress:]
		for len(d) > 0 {
			va := binary.LittleEndian.Uint32(d[0:4])
			d = d[4:]
			if va == 0 {
				break
			}
			if va&0x80000000 > 0 { // is Ordinal
				// TODO add dynimport ordinal support.
				//ord := va&0x0000FFFF
			} else {
				fn, _ := getString(names, int(va-ds.VirtualAddress+2))
				all = append(all, fn+":"+dt.dll)
			}
		}
	}

	return all, nil
}

// ImportedLibraries returns the names of all libraries
// referred to by the binary f that are expected to be
// linked with the binary at dynamic link time.
func (f *File) ImportedLibraries() ([]string, os.Error) {
	// TODO
	// cgo -dynimport don't use this for windows PE, so just return.
	return nil, nil
}
<commit_msg>debug/pe: fixes ImportedSymbols for Win64.<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 pe implements access to PE (Microsoft Windows Portable Executable) files.
package pe

import (
	"debug/dwarf"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"strconv"
)

// A File represents an open PE file.
type File struct {
	FileHeader
	Sections []*Section

	closer io.Closer
}

type SectionHeader struct {
	Name                 string
	VirtualSize          uint32
	VirtualAddress       uint32
	Size                 uint32
	Offset               uint32
	PointerToRelocations uint32
	PointerToLineNumbers uint32
	NumberOfRelocations  uint16
	NumberOfLineNumbers  uint16
	Characteristics      uint32
}

type Section struct {
	SectionHeader

	// Embed ReaderAt for ReadAt method.
	// Do not embed SectionReader directly
	// to avoid having Read and Seek.
	// If a client wants Read and Seek it must use
	// Open() to avoid fighting over the seek offset
	// with other clients.
	io.ReaderAt
	sr *io.SectionReader
}

type ImportDirectory struct {
	OriginalFirstThunk uint32
	TimeDateStamp      uint32
	ForwarderChain     uint32
	Name               uint32
	FirstThunk         uint32

	dll string
}

// Data reads and returns the contents of the PE section.
func (s *Section) Data() ([]byte, os.Error) {
	dat := make([]byte, s.sr.Size())
	n, err := s.sr.ReadAt(dat, 0)
	return dat[0:n], err
}

// Open returns a new ReadSeeker reading the PE section.
func (s *Section) Open() io.ReadSeeker { return io.NewSectionReader(s.sr, 0, 1<<63-1) }

type FormatError struct {
	off int64
	msg string
	val interface{}
}

func (e *FormatError) String() string {
	msg := e.msg
	if e.val != nil {
		msg += fmt.Sprintf(" '%v'", e.val)
	}
	msg += fmt.Sprintf(" in record at byte %#x", e.off)
	return msg
}

// Open opens the named file using os.Open and prepares it for use as a PE binary.
func Open(name string) (*File, os.Error) {
	f, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	ff, err := NewFile(f)
	if err != nil {
		f.Close()
		return nil, err
	}
	ff.closer = f
	return ff, nil
}

// Close closes the File.
// If the File was created using NewFile directly instead of Open,
// Close has no effect.
func (f *File) Close() os.Error {
	var err os.Error
	if f.closer != nil {
		err = f.closer.Close()
		f.closer = nil
	}
	return err
}

// NewFile creates a new File for accessing a PE binary in an underlying reader.
func NewFile(r io.ReaderAt) (*File, os.Error) {
	f := new(File)
	sr := io.NewSectionReader(r, 0, 1<<63-1)

	var dosheader [96]byte
	if _, err := r.ReadAt(dosheader[0:], 0); err != nil {
		return nil, err
	}
	var base int64
	if dosheader[0] == 'M' && dosheader[1] == 'Z' {
		var sign [4]byte
		r.ReadAt(sign[0:], int64(dosheader[0x3c]))
		if !(sign[0] == 'P' && sign[1] == 'E' && sign[2] == 0 && sign[3] == 0) {
			return nil, os.NewError("Invalid PE File Format.")
		}
		base = int64(dosheader[0x3c]) + 4
	} else {
		base = int64(0)
	}
	sr.Seek(base, os.SEEK_SET)
	if err := binary.Read(sr, binary.LittleEndian, &f.FileHeader); err != nil {
		return nil, err
	}
	if f.FileHeader.Machine != IMAGE_FILE_MACHINE_UNKNOWN && f.FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64 && f.FileHeader.Machine != IMAGE_FILE_MACHINE_I386 {
		return nil, os.NewError("Invalid PE File Format.")
	}
	// get symbol string table
	sr.Seek(int64(f.FileHeader.PointerToSymbolTable+18*f.FileHeader.NumberOfSymbols), os.SEEK_SET)
	var l uint32
	if err := binary.Read(sr, binary.LittleEndian, &l); err != nil {
		return nil, err
	}
	ss := make([]byte, l)
	if _, err := r.ReadAt(ss, int64(f.FileHeader.PointerToSymbolTable+18*f.FileHeader.NumberOfSymbols)); err != nil {
		return nil, err
	}
	sr.Seek(base, os.SEEK_SET)
	binary.Read(sr, binary.LittleEndian, &f.FileHeader)
	sr.Seek(int64(f.FileHeader.SizeOfOptionalHeader), os.SEEK_CUR) //Skip OptionalHeader
	f.Sections = make([]*Section, f.FileHeader.NumberOfSections)
	for i := 0; i < int(f.FileHeader.NumberOfSections); i++ {
		sh := new(SectionHeader32)
		if err := binary.Read(sr, binary.LittleEndian, sh); err != nil {
			return nil, err
		}
		var name string
		if sh.Name[0] == '\x2F' {
			si, _ := strconv.Atoi(cstring(sh.Name[1:]))
			name, _ = getString(ss, si)
		} else {
			name = cstring(sh.Name[0:])
		}
		s := new(Section)
		s.SectionHeader = SectionHeader{
			Name:                 name,
			VirtualSize:          uint32(sh.VirtualSize),
			VirtualAddress:       uint32(sh.VirtualAddress),
			Size:                 uint32(sh.SizeOfRawData),
			Offset:               uint32(sh.PointerToRawData),
			PointerToRelocations: uint32(sh.PointerToRelocations),
			PointerToLineNumbers: uint32(sh.PointerToLineNumbers),
			NumberOfRelocations:  uint16(sh.NumberOfRelocations),
			NumberOfLineNumbers:  uint16(sh.NumberOfLineNumbers),
			Characteristics:      uint32(sh.Characteristics),
		}
		s.sr = io.NewSectionReader(r, int64(s.SectionHeader.Offset), int64(s.SectionHeader.Size))
		s.ReaderAt = s.sr
		f.Sections[i] = s
	}
	return f, nil
}

func cstring(b []byte) string {
	var i int
	for i = 0; i < len(b) && b[i] != 0; i++ {
	}
	return string(b[0:i])
}

// getString extracts a string from symbol string table.
func getString(section []byte, start int) (string, bool) {
	if start < 0 || start >= len(section) {
		return "", false
	}

	for end := start; end < len(section); end++ {
		if section[end] == 0 {
			return string(section[start:end]), true
		}
	}
	return "", false
}

// Section returns the first section with the given name, or nil if no such
// section exists.
func (f *File) Section(name string) *Section {
	for _, s := range f.Sections {
		if s.Name == name {
			return s
		}
	}
	return nil
}

func (f *File) DWARF() (*dwarf.Data, os.Error) {
	// There are many other DWARF sections, but these
	// are the required ones, and the debug/dwarf package
	// does not use the others, so don't bother loading them.
	var names = [...]string{"abbrev", "info", "str"}
	var dat [len(names)][]byte
	for i, name := range names {
		name = ".debug_" + name
		s := f.Section(name)
		if s == nil {
			continue
		}
		b, err := s.Data()
		if err != nil && uint32(len(b)) < s.Size {
			return nil, err
		}
		dat[i] = b
	}

	abbrev, info, str := dat[0], dat[1], dat[2]
	return dwarf.New(abbrev, nil, nil, info, nil, nil, nil, str)
}

// ImportedSymbols returns the names of all symbols
// referred to by the binary f that are expected to be
// satisfied by other libraries at dynamic load time.
// It does not return weak symbols.
func (f *File) ImportedSymbols() ([]string, os.Error) {
	pe64 := f.Machine == IMAGE_FILE_MACHINE_AMD64
	ds := f.Section(".idata")
	if ds == nil {
		// not dynamic, so no libraries
		return nil, nil
	}
	d, err := ds.Data()
	if err != nil {
		return nil, err
	}
	var ida []ImportDirectory
	for len(d) > 0 {
		var dt ImportDirectory
		dt.OriginalFirstThunk = binary.LittleEndian.Uint32(d[0:4])
		dt.Name = binary.LittleEndian.Uint32(d[12:16])
		dt.FirstThunk = binary.LittleEndian.Uint32(d[16:20])
		d = d[20:]
		if dt.OriginalFirstThunk == 0 {
			break
		}
		ida = append(ida, dt)
	}
	names, _ := ds.Data()
	var all []string
	for _, dt := range ida {
		dt.dll, _ = getString(names, int(dt.Name-ds.VirtualAddress))
		d, _ = ds.Data()
		// seek to OriginalFirstThunk
		d = d[dt.OriginalFirstThunk-ds.VirtualAddress:]
		for len(d) > 0 {
			if pe64 { // 64bit
				va := binary.LittleEndian.Uint64(d[0:8])
				d = d[8:]
				if va == 0 {
					break
				}
				if va&0x8000000000000000 > 0 { // is Ordinal
					// TODO add dynimport ordinal support.
				} else {
					fn, _ := getString(names, int(uint32(va)-ds.VirtualAddress+2))
					all = append(all, fn+":"+dt.dll)
				}
			} else { // 32bit
				va := binary.LittleEndian.Uint32(d[0:4])
				d = d[4:]
				if va == 0 {
					break
				}
				if va&0x80000000 > 0 { // is Ordinal
					// TODO add dynimport ordinal support.
					//ord := va&0x0000FFFF
				} else {
					fn, _ := getString(names, int(va-ds.VirtualAddress+2))
					all = append(all, fn+":"+dt.dll)
				}
			}
		}
	}

	return all, nil
}

// ImportedLibraries returns the names of all libraries
// referred to by the binary f that are expected to be
// linked with the binary at dynamic link time.
func (f *File) ImportedLibraries() ([]string, os.Error) {
	// TODO
	// cgo -dynimport don't use this for windows PE, so just return.
	return nil, nil
}
<|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 build

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

// Path is a validated list of Trees derived from $GOPATH at init.
var Path []*Tree

// Tree describes a Go source tree, either $GOROOT or one from $GOPATH.
type Tree struct {
	Path   string
	Goroot bool
}

func newTree(p string) (*Tree, os.Error) {
	if !filepath.IsAbs(p) {
		return nil, os.NewError("must be absolute")
	}
	ep, err := filepath.EvalSymlinks(p)
	if err != nil {
		return nil, err
	}
	return &Tree{Path: ep}, nil
}

// SrcDir returns the tree's package source directory.
func (t *Tree) SrcDir() string {
	if t.Goroot {
		return filepath.Join(t.Path, "src", "pkg")
	}
	return filepath.Join(t.Path, "src")
}

// PkgDir returns the tree's package object directory.
func (t *Tree) PkgDir() string {
	goos, goarch := runtime.GOOS, runtime.GOARCH
	if e := os.Getenv("GOOS"); e != "" {
		goos = e
	}
	if e := os.Getenv("GOARCH"); e != "" {
		goarch = e
	}
	return filepath.Join(t.Path, "pkg", goos+"_"+goarch)
}

// BinDir returns the tree's binary executable directory.
func (t *Tree) BinDir() string {
	return filepath.Join(t.Path, "bin")
}

// HasSrc returns whether the given package's
// source can be found inside this Tree.
func (t *Tree) HasSrc(pkg string) bool {
	fi, err := os.Stat(filepath.Join(t.SrcDir(), pkg))
	if err != nil {
		return false
	}
	return fi.IsDirectory()
}

// HasPkg returns whether the given package's
// object file can be found inside this Tree.
func (t *Tree) HasPkg(pkg string) bool {
	fi, err := os.Stat(filepath.Join(t.PkgDir(), pkg+".a"))
	if err != nil {
		return false
	}
	return fi.IsRegular()
	// TODO(adg): check object version is consistent
}

var ErrNotFound = os.NewError("package could not be found locally")

// FindTree takes an import or filesystem path and returns the
// tree where the package source should be and the package import path.
func FindTree(path string) (tree *Tree, pkg string, err os.Error) {
	if isLocalPath(path) {
		if path, err = filepath.Abs(path); err != nil {
			return
		}
		if path, err = filepath.EvalSymlinks(path); err != nil {
			return
		}
		for _, t := range Path {
			tpath := t.SrcDir() + string(filepath.Separator)
			if !strings.HasPrefix(path, tpath) {
				continue
			}
			tree = t
			pkg = path[len(tpath):]
			return
		}
		err = fmt.Errorf("path %q not inside a GOPATH", path)
		return
	}
	tree = defaultTree
	pkg = path
	for _, t := range Path {
		if t.HasSrc(pkg) {
			tree = t
			return
		}
	}
	err = ErrNotFound
	return
}

// isLocalPath returns whether the given path is local (/foo ./foo ../foo . ..)
func isLocalPath(s string) bool {
	const sep = string(filepath.Separator)
	return strings.HasPrefix(s, sep) || strings.HasPrefix(s, "."+sep) || strings.HasPrefix(s, ".."+sep) || s == "." || s == ".."
}

var (
	// argument lists used by the build's gc and ld methods
	gcImportArgs []string
	ldImportArgs []string

	// default tree for remote packages
	defaultTree *Tree
)

// set up Path: parse and validate GOROOT and GOPATH variables
func init() {
	root := runtime.GOROOT()
	p, err := newTree(root)
	if err != nil {
		log.Fatalf("Invalid GOROOT %q: %v", root, err)
	}
	p.Goroot = true
	Path = []*Tree{p}

	for _, p := range filepath.SplitList(os.Getenv("GOPATH")) {
		if p == "" {
			continue
		}
		t, err := newTree(p)
		if err != nil {
			log.Printf("Invalid GOPATH %q: %v", p, err)
			continue
		}
		Path = append(Path, t)
		gcImportArgs = append(gcImportArgs, "-I", t.PkgDir())
		ldImportArgs = append(ldImportArgs, "-L", t.PkgDir())

		// select first GOPATH entry as default
		if defaultTree == nil {
			defaultTree = t
		}
	}

	// use GOROOT if no valid GOPATH specified
	if defaultTree == nil {
		defaultTree = Path[0]
	}
}
<commit_msg>go/build: less aggressive failure when GOROOT not found<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 build

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

// Path is a validated list of Trees derived from $GOROOT and $GOPATH at init.
var Path []*Tree

// Tree describes a Go source tree, either $GOROOT or one from $GOPATH.
type Tree struct {
	Path   string
	Goroot bool
}

func newTree(p string) (*Tree, os.Error) {
	if !filepath.IsAbs(p) {
		return nil, os.NewError("must be absolute")
	}
	ep, err := filepath.EvalSymlinks(p)
	if err != nil {
		return nil, err
	}
	return &Tree{Path: ep}, nil
}

// SrcDir returns the tree's package source directory.
func (t *Tree) SrcDir() string {
	if t.Goroot {
		return filepath.Join(t.Path, "src", "pkg")
	}
	return filepath.Join(t.Path, "src")
}

// PkgDir returns the tree's package object directory.
func (t *Tree) PkgDir() string {
	goos, goarch := runtime.GOOS, runtime.GOARCH
	if e := os.Getenv("GOOS"); e != "" {
		goos = e
	}
	if e := os.Getenv("GOARCH"); e != "" {
		goarch = e
	}
	return filepath.Join(t.Path, "pkg", goos+"_"+goarch)
}

// BinDir returns the tree's binary executable directory.
func (t *Tree) BinDir() string {
	return filepath.Join(t.Path, "bin")
}

// HasSrc returns whether the given package's
// source can be found inside this Tree.
func (t *Tree) HasSrc(pkg string) bool {
	fi, err := os.Stat(filepath.Join(t.SrcDir(), pkg))
	if err != nil {
		return false
	}
	return fi.IsDirectory()
}

// HasPkg returns whether the given package's
// object file can be found inside this Tree.
func (t *Tree) HasPkg(pkg string) bool {
	fi, err := os.Stat(filepath.Join(t.PkgDir(), pkg+".a"))
	if err != nil {
		return false
	}
	return fi.IsRegular()
	// TODO(adg): check object version is consistent
}

var (
	ErrNotFound     = os.NewError("go/build: package could not be found locally")
	ErrTreeNotFound = os.NewError("go/build: no valid GOROOT or GOPATH could be found")
)

// FindTree takes an import or filesystem path and returns the
// tree where the package source should be and the package import path.
func FindTree(path string) (tree *Tree, pkg string, err os.Error) {
	if isLocalPath(path) {
		if path, err = filepath.Abs(path); err != nil {
			return
		}
		if path, err = filepath.EvalSymlinks(path); err != nil {
			return
		}
		for _, t := range Path {
			tpath := t.SrcDir() + string(filepath.Separator)
			if !strings.HasPrefix(path, tpath) {
				continue
			}
			tree = t
			pkg = path[len(tpath):]
			return
		}
		err = fmt.Errorf("path %q not inside a GOPATH", path)
		return
	}
	tree = defaultTree
	pkg = path
	for _, t := range Path {
		if t.HasSrc(pkg) {
			tree = t
			return
		}
	}
	if tree == nil {
		err = ErrTreeNotFound
	} else {
		err = ErrNotFound
	}
	return
}

// isLocalPath returns whether the given path is local (/foo ./foo ../foo . ..)
func isLocalPath(s string) bool {
	const sep = string(filepath.Separator)
	return strings.HasPrefix(s, sep) || strings.HasPrefix(s, "."+sep) || strings.HasPrefix(s, ".."+sep) || s == "." || s == ".."
}

var (
	// argument lists used by the build's gc and ld methods
	gcImportArgs []string
	ldImportArgs []string

	// default tree for remote packages
	defaultTree *Tree
)

// set up Path: parse and validate GOROOT and GOPATH variables
func init() {
	root := runtime.GOROOT()
	t, err := newTree(root)
	if err != nil {
		log.Printf("go/build: invalid GOROOT %q: %v", root, err)
	} else {
		t.Goroot = true
		Path = []*Tree{t}
	}

	for _, p := range filepath.SplitList(os.Getenv("GOPATH")) {
		if p == "" {
			continue
		}
		t, err := newTree(p)
		if err != nil {
			log.Printf("go/build: invalid GOPATH %q: %v", p, err)
			continue
		}
		Path = append(Path, t)
		gcImportArgs = append(gcImportArgs, "-I", t.PkgDir())
		ldImportArgs = append(ldImportArgs, "-L", t.PkgDir())

		// select first GOPATH entry as default
		if defaultTree == nil {
			defaultTree = t
		}
	}

	// use GOROOT if no valid GOPATH specified
	if defaultTree == nil && len(Path) > 0 {
		defaultTree = Path[0]
	}
}
<|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 build

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

// Path is a validated list of Trees derived from $GOPATH at init.
var Path []*Tree

// Tree describes a Go source tree, either $GOROOT or one from $GOPATH.
type Tree struct {
	Path   string
	Goroot bool
}

func newTree(p string) (*Tree, os.Error) {
	if !filepath.IsAbs(p) {
		return nil, os.NewError("must be absolute")
	}
	ep, err := filepath.EvalSymlinks(p)
	if err != nil {
		return nil, err
	}
	return &Tree{Path: ep}, nil
}

// SrcDir returns the tree's package source directory.
func (t *Tree) SrcDir() string {
	if t.Goroot {
		return filepath.Join(t.Path, "src", "pkg")
	}
	return filepath.Join(t.Path, "src")
}

// PkgDir returns the tree's package object directory.
func (t *Tree) PkgDir() string {
	goos, goarch := runtime.GOOS, runtime.GOARCH
	if e := os.Getenv("GOOS"); e != "" {
		goos = e
	}
	if e := os.Getenv("GOARCH"); e != "" {
		goarch = e
	}
	return filepath.Join(t.Path, "pkg", goos+"_"+goarch)
}

// BinDir returns the tree's binary executable directory.
func (t *Tree) BinDir() string {
	return filepath.Join(t.Path, "bin")
}

// HasSrc returns whether the given package's
// source can be found inside this Tree.
func (t *Tree) HasSrc(pkg string) bool {
	fi, err := os.Stat(filepath.Join(t.SrcDir(), pkg))
	if err != nil {
		return false
	}
	return fi.IsDirectory()
}

// HasPkg returns whether the given package's
// object file can be found inside this Tree.
func (t *Tree) HasPkg(pkg string) bool {
	fi, err := os.Stat(filepath.Join(t.PkgDir(), pkg+".a"))
	if err != nil {
		return false
	}
	return fi.IsRegular()
	// TODO(adg): check object version is consistent
}

var ErrNotFound = os.NewError("package could not be found locally")

// FindTree takes an import or filesystem path and returns the
// tree where the package source should be and the package import path.
func FindTree(path string) (tree *Tree, pkg string, err os.Error) {
	if isLocalPath(path) {
		if path, err = filepath.Abs(path); err != nil {
			return
		}
		for _, t := range Path {
			tpath := t.SrcDir() + string(filepath.Separator)
			if !strings.HasPrefix(path, tpath) {
				continue
			}
			tree = t
			pkg = path[len(tpath):]
			return
		}
		err = fmt.Errorf("path %q not inside a GOPATH", path)
		return
	}
	tree = defaultTree
	pkg = path
	for _, t := range Path {
		if t.HasSrc(pkg) {
			tree = t
			return
		}
	}
	err = ErrNotFound
	return
}

// isLocalPath returns whether the given path is local (/foo ./foo ../foo . ..)
func isLocalPath(s string) bool {
	const sep = string(filepath.Separator)
	return strings.HasPrefix(s, sep) || strings.HasPrefix(s, "."+sep) || strings.HasPrefix(s, ".."+sep) || s == "." || s == ".."
}

var (
	// argument lists used by the build's gc and ld methods
	gcImportArgs []string
	ldImportArgs []string

	// default tree for remote packages
	defaultTree *Tree
)

// set up Path: parse and validate GOROOT and GOPATH variables
func init() {
	root := runtime.GOROOT()
	p, err := newTree(root)
	if err != nil {
		log.Fatalf("Invalid GOROOT %q: %v", root, err)
	}
	p.Goroot = true
	Path = []*Tree{p}

	for _, p := range filepath.SplitList(os.Getenv("GOPATH")) {
		if p == "" {
			continue
		}
		t, err := newTree(p)
		if err != nil {
			log.Printf("Invalid GOPATH %q: %v", p, err)
			continue
		}
		Path = append(Path, t)
		gcImportArgs = append(gcImportArgs, "-I", t.PkgDir())
		ldImportArgs = append(ldImportArgs, "-L", t.PkgDir())

		// select first GOPATH entry as default
		if defaultTree == nil {
			defaultTree = t
		}
	}

	// use GOROOT if no valid GOPATH specified
	if defaultTree == nil {
		defaultTree = Path[0]
	}
}
<commit_msg>go/build: evaluate symlinks before comparing path to GOPATH<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 build

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

// Path is a validated list of Trees derived from $GOPATH at init.
var Path []*Tree

// Tree describes a Go source tree, either $GOROOT or one from $GOPATH.
type Tree struct {
	Path   string
	Goroot bool
}

func newTree(p string) (*Tree, os.Error) {
	if !filepath.IsAbs(p) {
		return nil, os.NewError("must be absolute")
	}
	ep, err := filepath.EvalSymlinks(p)
	if err != nil {
		return nil, err
	}
	return &Tree{Path: ep}, nil
}

// SrcDir returns the tree's package source directory.
func (t *Tree) SrcDir() string {
	if t.Goroot {
		return filepath.Join(t.Path, "src", "pkg")
	}
	return filepath.Join(t.Path, "src")
}

// PkgDir returns the tree's package object directory.
func (t *Tree) PkgDir() string {
	goos, goarch := runtime.GOOS, runtime.GOARCH
	if e := os.Getenv("GOOS"); e != "" {
		goos = e
	}
	if e := os.Getenv("GOARCH"); e != "" {
		goarch = e
	}
	return filepath.Join(t.Path, "pkg", goos+"_"+goarch)
}

// BinDir returns the tree's binary executable directory.
func (t *Tree) BinDir() string {
	return filepath.Join(t.Path, "bin")
}

// HasSrc returns whether the given package's
// source can be found inside this Tree.
func (t *Tree) HasSrc(pkg string) bool {
	fi, err := os.Stat(filepath.Join(t.SrcDir(), pkg))
	if err != nil {
		return false
	}
	return fi.IsDirectory()
}

// HasPkg returns whether the given package's
// object file can be found inside this Tree.
func (t *Tree) HasPkg(pkg string) bool {
	fi, err := os.Stat(filepath.Join(t.PkgDir(), pkg+".a"))
	if err != nil {
		return false
	}
	return fi.IsRegular()
	// TODO(adg): check object version is consistent
}

var ErrNotFound = os.NewError("package could not be found locally")

// FindTree takes an import or filesystem path and returns the
// tree where the package source should be and the package import path.
func FindTree(path string) (tree *Tree, pkg string, err os.Error) {
	if isLocalPath(path) {
		if path, err = filepath.Abs(path); err != nil {
			return
		}
		if path, err = filepath.EvalSymlinks(path); err != nil {
			return
		}
		for _, t := range Path {
			tpath := t.SrcDir() + string(filepath.Separator)
			if !strings.HasPrefix(path, tpath) {
				continue
			}
			tree = t
			pkg = path[len(tpath):]
			return
		}
		err = fmt.Errorf("path %q not inside a GOPATH", path)
		return
	}
	tree = defaultTree
	pkg = path
	for _, t := range Path {
		if t.HasSrc(pkg) {
			tree = t
			return
		}
	}
	err = ErrNotFound
	return
}

// isLocalPath returns whether the given path is local (/foo ./foo ../foo . ..)
func isLocalPath(s string) bool {
	const sep = string(filepath.Separator)
	return strings.HasPrefix(s, sep) || strings.HasPrefix(s, "."+sep) || strings.HasPrefix(s, ".."+sep) || s == "." || s == ".."
}

var (
	// argument lists used by the build's gc and ld methods
	gcImportArgs []string
	ldImportArgs []string

	// default tree for remote packages
	defaultTree *Tree
)

// set up Path: parse and validate GOROOT and GOPATH variables
func init() {
	root := runtime.GOROOT()
	p, err := newTree(root)
	if err != nil {
		log.Fatalf("Invalid GOROOT %q: %v", root, err)
	}
	p.Goroot = true
	Path = []*Tree{p}

	for _, p := range filepath.SplitList(os.Getenv("GOPATH")) {
		if p == "" {
			continue
		}
		t, err := newTree(p)
		if err != nil {
			log.Printf("Invalid GOPATH %q: %v", p, err)
			continue
		}
		Path = append(Path, t)
		gcImportArgs = append(gcImportArgs, "-I", t.PkgDir())
		ldImportArgs = append(ldImportArgs, "-L", t.PkgDir())

		// select first GOPATH entry as default
		if defaultTree == nil {
			defaultTree = t
		}
	}

	// use GOROOT if no valid GOPATH specified
	if defaultTree == nil {
		defaultTree = Path[0]
	}
}
<|endoftext|>


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

package mq

import (
	"context"
	"fmt"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/log"
	"github.com/shawnfeng/sutil/slog/slog"
	"github.com/shawnfeng/sutil/stime"
	"time"
	// kafka "github.com/segmentio/kafka-go"
)

const (
	spanLogKeyKey       = "key"
	spanLogKeyTopic     = "topic"
	spanLogKeyGroupId   = "groupId"
	spanLogKeyPartition = "partition"
)

var mqOpDurationLimit = 10 * time.Millisecond

type Message struct {
	Key   string
	Value interface{}
}

func WriteMsg(ctx context.Context, topic string, key string, value interface{}) error {
	fun := "WriteMsg -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.WriteMsg")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyTopic, topic),
		log.String(spanLogKeyKey, key))

	//todo flag
	writer := DefaultInstanceManager.getWriter("", ROLE_TYPE_WRITER, topic, "", 0)
	if writer == nil {
		slog.Errorf(ctx, "%s getWriter err, topic: %s", fun, topic)
		return fmt.Errorf("%s, getWriter err, topic: %s", fun, topic)
	}

	payload, err := generatePayload(ctx, value)
	if err != nil {
		slog.Errorf(ctx, "%s generatePayload err, topic: %s", fun, topic)
		return fmt.Errorf("%s, generatePayload err, topic: %s", fun, topic)
	}

	st := stime.NewTimeStat()
	defer func() {
		dur := st.Duration()
		if dur > mqOpDurationLimit {
			slog.Infof(ctx, "%s topic:%s dur:%d", fun, topic, dur)
		}
	}()

	return writer.WriteMsg(ctx, key, payload)
}

func WriteMsgs(ctx context.Context, topic string, msgs ...Message) error {
	fun := "WriteMsgs -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.WriteMsgs")
	defer span.Finish()
	span.LogFields(log.String(spanLogKeyTopic, topic))

	//todo flag
	writer := DefaultInstanceManager.getWriter("", ROLE_TYPE_WRITER, topic, "", 0)
	if writer == nil {
		slog.Errorf(ctx, "%s getWriter err, topic: %s", fun, topic)
		return fmt.Errorf("%s, getWriter err, topic: %s", fun, topic)
	}

	nmsgs, err := generateMsgsPayload(ctx, msgs...)
	if err != nil {
		slog.Errorf(ctx, "%s generateMsgsPayload err, topic: %s", fun, topic)
		return fmt.Errorf("%s, generateMsgsPayload err, topic: %s", fun, topic)
	}

	st := stime.NewTimeStat()
	defer func() {
		dur := st.Duration()
		if dur > mqOpDurationLimit {
			slog.Infof(ctx, "%s topic:%s dur:%d", fun, topic, dur)
		}
	}()

	return writer.WriteMsgs(ctx, nmsgs...)
}

// 读完消息后会自动提交offset
func ReadMsgByGroup(ctx context.Context, topic, groupId string, value interface{}) (context.Context, error) {
	fun := "ReadMsgByGroup -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.ReadMsgByGroup")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyTopic, topic),
		log.String(spanLogKeyGroupId, groupId))

	//todo flag
	reader := DefaultInstanceManager.getReader("", ROLE_TYPE_READER, topic, groupId, 0)
	if reader == nil {
		slog.Errorf(ctx, "%s getReader err, topic: %s", fun, topic)
		return nil, fmt.Errorf("%s, getReader err, topic: %s", fun, topic)
	}

	var payload Payload
	st := stime.NewTimeStat()

	err := reader.ReadMsg(ctx, &payload, value)

	dur := st.Duration()
	if dur > mqOpDurationLimit {
		slog.Infof(ctx, "%s topic:%s groupId:%s dur:%d", fun, topic, groupId, dur)
	}

	if err != nil {
		slog.Errorf(ctx, "%s ReadMsg err, topic: %s", fun, topic)
		return nil, fmt.Errorf("%s, ReadMsg err, topic: %s", fun, topic)
	}

	if len(payload.Value) == 0 {
		return context.TODO(), nil
	}

	mctx, err := parsePayload(&payload, "mq.ReadMsgByGroup", value)
	mspan := opentracing.SpanFromContext(mctx)
	if mspan != nil {
		defer mspan.Finish()
		mspan.LogFields(
			log.String(spanLogKeyTopic, topic),
			log.String(spanLogKeyGroupId, groupId))
	}
	return mctx, err
}

//
func ReadMsgByPartition(ctx context.Context, topic string, partition int, value interface{}) (context.Context, error) {
	fun := "ReadMsgByPartition -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.ReadMsgByPartition")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyTopic, topic),
		log.Int(spanLogKeyPartition, partition))

	//todo flag
	reader := DefaultInstanceManager.getReader("", ROLE_TYPE_READER, topic, "", partition)
	if reader == nil {
		slog.Errorf(ctx, "%s getReader err, topic: %s", fun, topic)
		return nil, fmt.Errorf("%s, getReader err, topic: %s", fun, topic)
	}

	var payload Payload
	st := stime.NewTimeStat()

	err := reader.ReadMsg(ctx, &payload, value)

	dur := st.Duration()
	if dur > mqOpDurationLimit {
		slog.Infof(ctx, "%s topic:%s partition:%d dur:%d", fun, topic, partition, dur)
	}

	if err != nil {
		slog.Errorf(ctx, "%s ReadMsg err, topic: %s", fun, topic)
		return nil, fmt.Errorf("%s, ReadMsg err, topic: %s", fun, topic)
	}

	if len(payload.Value) == 0 {
		return context.TODO(), nil
	}

	mctx, err := parsePayload(&payload, "mq.ReadMsgByPartition", value)
	mspan := opentracing.SpanFromContext(mctx)
	if mspan != nil {
		defer mspan.Finish()
		mspan.LogFields(
			log.String(spanLogKeyTopic, topic),
			log.Int(spanLogKeyPartition, partition))
	}
	return mctx, err
}

// 读完消息后不会自动提交offset,需要手动调用Handle.CommitMsg方法来提交offset
func FetchMsgByGroup(ctx context.Context, topic, groupId string, value interface{}) (context.Context, Handler, error) {
	fun := "FetchMsgByGroup -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.FetchMsgByGroup")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyTopic, topic),
		log.String(spanLogKeyGroupId, groupId))

	//todo flag
	reader := DefaultInstanceManager.getReader("", ROLE_TYPE_READER, topic, groupId, 0)
	if reader == nil {
		slog.Errorf(ctx, "%s getReader err, topic: %s", fun, topic)
		return nil, nil, fmt.Errorf("%s, getReader err, topic: %s", fun, topic)
	}

	var payload Payload
	st := stime.NewTimeStat()

	handler, err := reader.FetchMsg(ctx, &payload, value)

	dur := st.Duration()
	if dur > mqOpDurationLimit {
		slog.Infof(ctx, "%s topic:%s groupId:%s dur:%d", fun, topic, groupId, dur)
	}

	if err != nil {
		slog.Errorf(ctx, "%s ReadMsg err, topic: %s", fun, topic)
		return nil, nil, fmt.Errorf("%s, ReadMsg err, topic: %s", fun, topic)
	}

	if len(payload.Value) == 0 {
		return context.TODO(), handler, nil
	}

	mctx, err := parsePayload(&payload, "mq.FetchMsgByGroup", value)
	mspan := opentracing.SpanFromContext(mctx)
	if mspan != nil {
		defer mspan.Finish()
		mspan.LogFields(
			log.String(spanLogKeyTopic, topic),
			log.String(spanLogKeyGroupId, groupId))
	}
	return mctx, handler, err
}

func Close() {
	DefaultInstanceManager.Close()
}
<commit_msg>add slow keyword to log<commit_after>// Copyright 2014 The mqrouter Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mq

import (
	"context"
	"fmt"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/log"
	"github.com/shawnfeng/sutil/slog/slog"
	"github.com/shawnfeng/sutil/stime"
	"time"
	// kafka "github.com/segmentio/kafka-go"
)

const (
	spanLogKeyKey       = "key"
	spanLogKeyTopic     = "topic"
	spanLogKeyGroupId   = "groupId"
	spanLogKeyPartition = "partition"
)

var mqOpDurationLimit = 10 * time.Millisecond

type Message struct {
	Key   string
	Value interface{}
}

func WriteMsg(ctx context.Context, topic string, key string, value interface{}) error {
	fun := "WriteMsg -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.WriteMsg")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyTopic, topic),
		log.String(spanLogKeyKey, key))

	//todo flag
	writer := DefaultInstanceManager.getWriter("", ROLE_TYPE_WRITER, topic, "", 0)
	if writer == nil {
		slog.Errorf(ctx, "%s getWriter err, topic: %s", fun, topic)
		return fmt.Errorf("%s, getWriter err, topic: %s", fun, topic)
	}

	payload, err := generatePayload(ctx, value)
	if err != nil {
		slog.Errorf(ctx, "%s generatePayload err, topic: %s", fun, topic)
		return fmt.Errorf("%s, generatePayload err, topic: %s", fun, topic)
	}

	st := stime.NewTimeStat()
	defer func() {
		dur := st.Duration()
		if dur > mqOpDurationLimit {
			slog.Infof(ctx, "%s slow topic:%s dur:%d", fun, topic, dur)
		}
	}()

	return writer.WriteMsg(ctx, key, payload)
}

func WriteMsgs(ctx context.Context, topic string, msgs ...Message) error {
	fun := "WriteMsgs -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.WriteMsgs")
	defer span.Finish()
	span.LogFields(log.String(spanLogKeyTopic, topic))

	//todo flag
	writer := DefaultInstanceManager.getWriter("", ROLE_TYPE_WRITER, topic, "", 0)
	if writer == nil {
		slog.Errorf(ctx, "%s getWriter err, topic: %s", fun, topic)
		return fmt.Errorf("%s, getWriter err, topic: %s", fun, topic)
	}

	nmsgs, err := generateMsgsPayload(ctx, msgs...)
	if err != nil {
		slog.Errorf(ctx, "%s generateMsgsPayload err, topic: %s", fun, topic)
		return fmt.Errorf("%s, generateMsgsPayload err, topic: %s", fun, topic)
	}

	st := stime.NewTimeStat()
	defer func() {
		dur := st.Duration()
		if dur > mqOpDurationLimit {
			slog.Infof(ctx, "%s slow topic:%s dur:%d", fun, topic, dur)
		}
	}()

	return writer.WriteMsgs(ctx, nmsgs...)
}

// 读完消息后会自动提交offset
func ReadMsgByGroup(ctx context.Context, topic, groupId string, value interface{}) (context.Context, error) {
	fun := "ReadMsgByGroup -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.ReadMsgByGroup")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyTopic, topic),
		log.String(spanLogKeyGroupId, groupId))

	//todo flag
	reader := DefaultInstanceManager.getReader("", ROLE_TYPE_READER, topic, groupId, 0)
	if reader == nil {
		slog.Errorf(ctx, "%s getReader err, topic: %s", fun, topic)
		return nil, fmt.Errorf("%s, getReader err, topic: %s", fun, topic)
	}

	var payload Payload
	st := stime.NewTimeStat()

	err := reader.ReadMsg(ctx, &payload, value)

	dur := st.Duration()
	if dur > mqOpDurationLimit {
		slog.Infof(ctx, "%s slow topic:%s groupId:%s dur:%d", fun, topic, groupId, dur)
	}

	if err != nil {
		slog.Errorf(ctx, "%s ReadMsg err, topic: %s", fun, topic)
		return nil, fmt.Errorf("%s, ReadMsg err, topic: %s", fun, topic)
	}

	if len(payload.Value) == 0 {
		return context.TODO(), nil
	}

	mctx, err := parsePayload(&payload, "mq.ReadMsgByGroup", value)
	mspan := opentracing.SpanFromContext(mctx)
	if mspan != nil {
		defer mspan.Finish()
		mspan.LogFields(
			log.String(spanLogKeyTopic, topic),
			log.String(spanLogKeyGroupId, groupId))
	}
	return mctx, err
}

//
func ReadMsgByPartition(ctx context.Context, topic string, partition int, value interface{}) (context.Context, error) {
	fun := "ReadMsgByPartition -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.ReadMsgByPartition")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyTopic, topic),
		log.Int(spanLogKeyPartition, partition))

	//todo flag
	reader := DefaultInstanceManager.getReader("", ROLE_TYPE_READER, topic, "", partition)
	if reader == nil {
		slog.Errorf(ctx, "%s getReader err, topic: %s", fun, topic)
		return nil, fmt.Errorf("%s, getReader err, topic: %s", fun, topic)
	}

	var payload Payload
	st := stime.NewTimeStat()

	err := reader.ReadMsg(ctx, &payload, value)

	dur := st.Duration()
	if dur > mqOpDurationLimit {
		slog.Infof(ctx, "%s slow topic:%s partition:%d dur:%d", fun, topic, partition, dur)
	}

	if err != nil {
		slog.Errorf(ctx, "%s ReadMsg err, topic: %s", fun, topic)
		return nil, fmt.Errorf("%s, ReadMsg err, topic: %s", fun, topic)
	}

	if len(payload.Value) == 0 {
		return context.TODO(), nil
	}

	mctx, err := parsePayload(&payload, "mq.ReadMsgByPartition", value)
	mspan := opentracing.SpanFromContext(mctx)
	if mspan != nil {
		defer mspan.Finish()
		mspan.LogFields(
			log.String(spanLogKeyTopic, topic),
			log.Int(spanLogKeyPartition, partition))
	}
	return mctx, err
}

// 读完消息后不会自动提交offset,需要手动调用Handle.CommitMsg方法来提交offset
func FetchMsgByGroup(ctx context.Context, topic, groupId string, value interface{}) (context.Context, Handler, error) {
	fun := "FetchMsgByGroup -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "mq.FetchMsgByGroup")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyTopic, topic),
		log.String(spanLogKeyGroupId, groupId))

	//todo flag
	reader := DefaultInstanceManager.getReader("", ROLE_TYPE_READER, topic, groupId, 0)
	if reader == nil {
		slog.Errorf(ctx, "%s getReader err, topic: %s", fun, topic)
		return nil, nil, fmt.Errorf("%s, getReader err, topic: %s", fun, topic)
	}

	var payload Payload
	st := stime.NewTimeStat()

	handler, err := reader.FetchMsg(ctx, &payload, value)

	dur := st.Duration()
	if dur > mqOpDurationLimit {
		slog.Infof(ctx, "%s slow topic:%s groupId:%s dur:%d", fun, topic, groupId, dur)
	}

	if err != nil {
		slog.Errorf(ctx, "%s ReadMsg err, topic: %s", fun, topic)
		return nil, nil, fmt.Errorf("%s, ReadMsg err, topic: %s", fun, topic)
	}

	if len(payload.Value) == 0 {
		return context.TODO(), handler, nil
	}

	mctx, err := parsePayload(&payload, "mq.FetchMsgByGroup", value)
	mspan := opentracing.SpanFromContext(mctx)
	if mspan != nil {
		defer mspan.Finish()
		mspan.LogFields(
			log.String(spanLogKeyTopic, topic),
			log.String(spanLogKeyGroupId, groupId))
	}
	return mctx, handler, err
}

func Close() {
	DefaultInstanceManager.Close()
}
<|endoftext|>


<commit_before>package entry

import (
	"bytes"
	"fmt"
	urlparse "net/url"
	"regexp"
	"strings"

	libm "bitbucket.org/ikeikeikeike/antenna/lib/models"
	gq "github.com/PuerkitoBio/goquery"
	behavior "github.com/ikeikeikeike/gopkg/net/http"
	"github.com/ikeikeikeike/gopkg/str"

	"golang.org/x/net/html"

	"bitbucket.org/ikeikeikeike/antenna/models"
)

var (
	defaultUserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_5) AppleWebKit/600.3.18 (KHTML, like Gecko) Version/7.1.3 Safari/537.85.12"

	embedUrls = []string{
		`flashservice.xvideos.com/embedframe`,
		`video.fc2.com/content`,
		`video.fc2.com/a/content`,
		`video.fc2.com/ja/a/content`,
		`www.tokyo-tube.com/embedcode`,
		`xhamster.com/xembed.php`,
		`www.pornhost.com/embed`,
		`www.tube8.com/embed/asian`,
		`www.pipii.tv/player`,
		`jp.pornhub.com/embed`,
		`www.pornhub.com/embed`,
		`japan-whores.com/embed`,
		`www.tokyo-porn-tube.com/embedcode`,
	}

	// TODO: will be support future below.

	//

	// <object classid="clsid:d27cdb6e-ae6d-11cf-96b8-444553540000" codebase="http://fpdownload.macromedia.com/pub/shockwave/cabs/flash/swflash.cab#version=8,0,0,0" wmode="transparent" width="448" height="284" id="flv2" align="middle"><param name="allowScriptAccess" value="sameDomain" /><param name="movie" value="http://video.fc2.com/flv2.swf?i=20150206aZzPs9Mg&d=3063&movie_stop=off&no_progressive=1&otag=1&sj=28000&rel=1" /><param name="quality" value="high" /><param name="bgcolor" value="#ffffff" /><param name="allowFullScreen" value="true" /><embed src="http://video.fc2.com/flv2.swf?i=20150206aZzPs9Mg&d=3063&movie_stop=off&no_progressive=1&otag=1&sj=28000&rel=1" quality="high" bgcolor="#ffffff" wmode="transparent" width="448" height="284" name="flv2" align="middle" allowScriptAccess="sameDomain" type="application/x-shockwave-flash" pluginspage="http://www.macromedia.com/go/getflashplayer" allowFullScreen="true" /></object>
<a href="http://video.fc2.com/content/20150206aZzPs9Mg/" title="対応" rel="nofollow" 【、可愛いが好き

	videoHrefs = []string{
		`www.xvideos.com/video`,
		`jp.xvideos.com/video`,
		`video.fc2.com/content`,
		`video.fc2.com/a/content`,
		`video.fc2.com/ja/a/content`,
		`www.tokyo-tube.com/video`,
		`ero-video.net/movie`,
		`jp.xhamster.com/movies`,
		`www.xhamster.com/movies`,
		`asg.to/contentsPage.html`,
		`www.pornhost.com`,
		`jp.pornhost.com`,
		`www.tube8.com`,
		`www.redtube.com`,
		`www.pornhub.com/view_video.php`,
		`jp.pornhub.com/view_video.php`,
		`japan-whores.com/videos`,
		`www.tokyo-porn-tube.com/video`,
		`www.flashx.tv`,
		// `www.thisav.com/video`,  // unko
	}
)

type Extractor struct {
	*behavior.UserBehavior

	entry *models.Entry
	doc   *gq.Document
}

func NewExtractor() *Extractor {
	return &Extractor{
		UserBehavior: behavior.NewUserBehavior(),
	}
}

func (e *Extractor) Doc(urlStr string) (*gq.Document, error) {
	resp, err := e.Behave(urlStr)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return gq.NewDocumentFromResponse(resp)
}

/*
	Do parsing document
*/
func (e *Extractor) Do(entry *models.Entry) error {
	doc, err := e.Doc(entry.Url)
	if err != nil {
		return err
	}

	e.doc = doc
	e.entry = entry
	return nil
}

/*
	Extracting ExternalURLs
*/
func (e *Extractor) Urls() (urls []string) {
	var re = regexp.MustCompile(strings.Join(videoHrefs, `|`))

	e.doc.Find("a").Each(func(i int, s *gq.Selection) {
		h, ok := s.Attr("href")
		if !ok {
			return
		}
		u, err := urlparse.Parse(h)
		if err != nil || u.Host == "" {
			return
		}

		if re.MatchString(h) {
			urls = append(urls, h)
		}
	})

	return
}

// iframe or script or object
// srcとurlでcontains ableなもの
func (e *Extractor) Codes() (codes []string) {

	var render = func(sel *gq.Selection) {
		var buf bytes.Buffer
		sel.Each(func(i int, s *gq.Selection) {
			html.Render(&buf, s.Nodes[0])
			codes = append(codes, str.Clean(buf.String()))
		})
	}

	// sel := e.doc.Find("iframe,script,embed,object").FilterFunction(func(i int, s *gq.Selection) bool {
	var re = regexp.MustCompile(strings.Join(embedUrls, `|`))
	e.doc.Find("iframe,script").Each(func(i int, s *gq.Selection) {
		src, ok := s.Attr("src")
		if ok && re.MatchString(src) {
			render(s)
		}
		url, ok := s.Attr("url")
		if ok && re.MatchString(url) {
			render(s)
		}
		if strings.Contains(src, "ero-video.net/js/embed_evplayer") {
			render(s)
			render(s.Next())
		}
		if strings.Contains(src, "asg.to/js/past_uraui") {
			render(s)
			render(s.Next())
		}
	})

	if len(codes) <= 0 {
		e.doc.Find("object").Each(func(i int, s *gq.Selection) {
			data, ok := s.Attr("data")
			if ok && strings.Contains(data, "embed.redtube.com/player") {
				render(s)
			}
		})
	}

	return
}

/*
	Extracting by string array in Title, Content

	.. note:: 3 chars less word will not extract in function.
*/
func (e *Extractor) ByNames(names []string) (choices []string) {
	for _, name := range names {
		if len([]rune(name)) > 2 && !libm.ReHK3.MatchString(name) {
			if strings.Contains(e.entry.Title, name) {
				choices = append(choices, name)
			} else if strings.Contains(e.entry.Content, name) {
				choices = append(choices, name)
			} else {
				for _, t := range e.entry.Tags {
					if strings.Contains(t.Name, name) {
						choices = append(choices, name)
					}
				}
			}
		}
	}
	return
}

func (e *Extractor) Duration() {}

type Img struct {
	Src string
	Alt string
}

/*
	Extracting img tag
*/
func (e *Extractor) Imgs() (imgs []*Img) {
	u, err := urlparse.Parse(strings.TrimSpace(e.entry.Url))
	if err != nil || u.Host == "" {
		return
	}

	var src string
	domain := strings.Split(u.Host, ":")[0]

	if strings.Contains(domain, "livedoor.jp") {
		src = strings.Split(u.Path, "/")[1]
	} else if strings.Contains(domain, "fc2.com") {
		src = strings.Split(domain, ".")[0]
	} else {
		parts := strings.Split(domain, ".")
		src = parts[len(parts)-2]
	}

	selector := fmt.Sprintf(`img[src*='%s']`, src)
	e.doc.Find(selector).Each(func(i int, s *gq.Selection) {
		src, ok := s.Attr("src")
		if !ok {
			return
		}
		alt, _ := s.Attr("alt")
		if alt == "" {
			alt, _ = s.Attr("title")
		}
		imgs = append(imgs, &Img{Src: src, Alt: alt})
	})

	return
}
<commit_msg>support multiple extract<commit_after>package entry

import (
	"bytes"
	"fmt"
	urlparse "net/url"
	"regexp"
	"strings"

	libm "bitbucket.org/ikeikeikeike/antenna/lib/models"
	gq "github.com/PuerkitoBio/goquery"
	behavior "github.com/ikeikeikeike/gopkg/net/http"
	"github.com/ikeikeikeike/gopkg/str"

	"golang.org/x/net/html"

	"bitbucket.org/ikeikeikeike/antenna/models"
)

var (
	defaultUserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_5) AppleWebKit/600.3.18 (KHTML, like Gecko) Version/7.1.3 Safari/537.85.12"

	embedUrls = []string{
		`flashservice.xvideos.com/embedframe`,
		`video.fc2.com/content`,
		`video.fc2.com/a/content`,
		`video.fc2.com/ja/a/content`,
		`www.tokyo-tube.com/embedcode`,
		`xhamster.com/xembed.php`,
		`www.pornhost.com/embed`,
		`www.tube8.com/embed/asian`,
		`www.pipii.tv/player`,
		`jp.pornhub.com/embed`,
		`www.pornhub.com/embed`,
		`japan-whores.com/embed`,
		`www.tokyo-porn-tube.com/embedcode`,
	}

	// TODO: will be support future below.

	//

	// <object classid="clsid:d27cdb6e-ae6d-11cf-96b8-444553540000" codebase="http://fpdownload.macromedia.com/pub/shockwave/cabs/flash/swflash.cab#version=8,0,0,0" wmode="transparent" width="448" height="284" id="flv2" align="middle"><param name="allowScriptAccess" value="sameDomain" /><param name="movie" value="http://video.fc2.com/flv2.swf?i=20150206aZzPs9Mg&d=3063&movie_stop=off&no_progressive=1&otag=1&sj=28000&rel=1" /><param name="quality" value="high" /><param name="bgcolor" value="#ffffff" /><param name="allowFullScreen" value="true" /><embed src="http://video.fc2.com/flv2.swf?i=20150206aZzPs9Mg&d=3063&movie_stop=off&no_progressive=1&otag=1&sj=28000&rel=1" quality="high" bgcolor="#ffffff" wmode="transparent" width="448" height="284" name="flv2" align="middle" allowScriptAccess="sameDomain" type="application/x-shockwave-flash" pluginspage="http://www.macromedia.com/go/getflashplayer" allowFullScreen="true" /></object>
<a href="http://video.fc2.com/content/20150206aZzPs9Mg/" title="対応" rel="nofollow" 【、可愛いが好き

	videoHrefs = []string{
		`www.xvideos.com/video`,
		`jp.xvideos.com/video`,
		`video.fc2.com/content`,
		`video.fc2.com/a/content`,
		`video.fc2.com/ja/a/content`,
		`www.tokyo-tube.com/video`,
		`ero-video.net/movie`,
		`jp.xhamster.com/movies`,
		`www.xhamster.com/movies`,
		`asg.to/contentsPage.html`,
		`www.pornhost.com`,
		`jp.pornhost.com`,
		`www.tube8.com`,
		`www.redtube.com`,
		`www.pornhub.com/view_video.php`,
		`jp.pornhub.com/view_video.php`,
		`japan-whores.com/videos`,
		`www.tokyo-porn-tube.com/video`,
		`www.flashx.tv`,
		// `www.thisav.com/video`,  // unko
	}
)

type Extractor struct {
	*behavior.UserBehavior

	entry *models.Entry
	doc   *gq.Document
}

func NewExtractor() *Extractor {
	return &Extractor{
		UserBehavior: behavior.NewUserBehavior(),
	}
}

func (e *Extractor) Doc(urlStr string) (*gq.Document, error) {
	resp, err := e.Behave(urlStr)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return gq.NewDocumentFromResponse(resp)
}

/*
	Do parsing document
*/
func (e *Extractor) Do(entry *models.Entry) error {
	doc, err := e.Doc(entry.Url)
	if err != nil {
		return err
	}

	e.doc = doc
	e.entry = entry
	return nil
}

/*
	Extracting ExternalURLs
*/
func (e *Extractor) Urls() (urls []string) {
	var re = regexp.MustCompile(strings.Join(videoHrefs, `|`))

	e.doc.Find("a").Each(func(i int, s *gq.Selection) {
		h, ok := s.Attr("href")
		if !ok {
			return
		}
		u, err := urlparse.Parse(h)
		if err != nil || u.Host == "" {
			return
		}

		if re.MatchString(h) {
			urls = append(urls, h)
		}
	})

	return
}

// iframe or script or object
// srcとurlでcontains ableなもの
func (e *Extractor) Codes() (codes []string) {

	var render = func(selectors ...*gq.Selection) {
		var raw string
		for _, sel := range selectors {
			var buf bytes.Buffer
			sel.Each(func(i int, s *gq.Selection) {
				html.Render(&buf, s.Nodes[0])
				raw += str.Clean(buf.String())
			})
		}
		codes = append(codes, raw)
	}

	// sel := e.doc.Find("iframe,script,embed,object").FilterFunction(func(i int, s *gq.Selection) bool {
	var re = regexp.MustCompile(strings.Join(embedUrls, `|`))
	e.doc.Find("iframe,script").Each(func(i int, s *gq.Selection) {
		src, ok := s.Attr("src")
		if ok && re.MatchString(src) {
			render(s)
		}
		url, ok := s.Attr("url")
		if ok && re.MatchString(url) {
			render(s)
		}
		if strings.Contains(src, "ero-video.net/js/embed_evplayer") {
			render(s, s.Next())
		}
		if strings.Contains(src, "asg.to/js/past_uraui") {
			render(s, s.Next())
		}
	})

	if len(codes) <= 0 {
		e.doc.Find("object").Each(func(i int, s *gq.Selection) {
			data, ok := s.Attr("data")
			if ok && strings.Contains(data, "embed.redtube.com/player") {
				render(s)
			}
		})
	}

	return
}

/*
	Extracting by string array in Title, Content

	.. note:: 3 chars less word will not extract in function.
*/
func (e *Extractor) ByNames(names []string) (choices []string) {
	for _, name := range names {
		if len([]rune(name)) > 2 && !libm.ReHK3.MatchString(name) {
			if strings.Contains(e.entry.Title, name) {
				choices = append(choices, name)
			} else if strings.Contains(e.entry.Content, name) {
				choices = append(choices, name)
			} else {
				for _, t := range e.entry.Tags {
					if strings.Contains(t.Name, name) {
						choices = append(choices, name)
					}
				}
			}
		}
	}
	return
}

func (e *Extractor) Duration() {}

type Img struct {
	Src string
	Alt string
}

/*
	Extracting img tag
*/
func (e *Extractor) Imgs() (imgs []*Img) {
	u, err := urlparse.Parse(strings.TrimSpace(e.entry.Url))
	if err != nil || u.Host == "" {
		return
	}

	var src string
	domain := strings.Split(u.Host, ":")[0]

	if strings.Contains(domain, "livedoor.jp") {
		src = strings.Split(u.Path, "/")[1]
	} else if strings.Contains(domain, "fc2.com") {
		src = strings.Split(domain, ".")[0]
	} else {
		parts := strings.Split(domain, ".")
		src = parts[len(parts)-2]
	}

	selector := fmt.Sprintf(`img[src*='%s']`, src)
	e.doc.Find(selector).Each(func(i int, s *gq.Selection) {
		src, ok := s.Attr("src")
		if !ok {
			return
		}
		alt, _ := s.Attr("alt")
		if alt == "" {
			alt, _ = s.Attr("title")
		}
		imgs = append(imgs, &Img{Src: src, Alt: alt})
	})

	return
}
<|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 net

import (
	"math/rand"
	"sort"
)

// DNSError represents a DNS lookup error.
type DNSError struct {
	Err       string // description of the error
	Name      string // name looked for
	Server    string // server used
	IsTimeout bool
}

func (e *DNSError) Error() string {
	if e == nil {
		return "<nil>"
	}
	s := "lookup " + e.Name
	if e.Server != "" {
		s += " on " + e.Server
	}
	s += ": " + e.Err
	return s
}

func (e *DNSError) Timeout() bool   { return e.IsTimeout }
func (e *DNSError) Temporary() bool { return e.IsTimeout }

const noSuchHost = "no such host"

// reverseaddr returns the in-addr.arpa. or ip6.arpa. hostname of the IP
// address addr suitable for rDNS (PTR) record lookup or an error if it fails
// to parse the IP address.
func reverseaddr(addr string) (arpa string, err error) {
	ip := ParseIP(addr)
	if ip == nil {
		return "", &DNSError{Err: "unrecognized address", Name: addr}
	}
	if ip.To4() != nil {
		return itoa(int(ip[15])) + "." + itoa(int(ip[14])) + "." + itoa(int(ip[13])) + "." +
			itoa(int(ip[12])) + ".in-addr.arpa.", nil
	}
	// Must be IPv6
	buf := make([]byte, 0, len(ip)*4+len("ip6.arpa."))
	// Add it, in reverse, to the buffer
	for i := len(ip) - 1; i >= 0; i-- {
		v := ip[i]
		buf = append(buf, hexDigit[v&0xF])
		buf = append(buf, '.')
		buf = append(buf, hexDigit[v>>4])
		buf = append(buf, '.')
	}
	// Append "ip6.arpa." and return (buf already has the final .)
	buf = append(buf, "ip6.arpa."...)
	return string(buf), nil
}

// Find answer for name in dns message.
// On return, if err == nil, addrs != nil.
func answer(name, server string, dns *dnsMsg, qtype uint16) (cname string, addrs []dnsRR, err error) {
	addrs = make([]dnsRR, 0, len(dns.answer))

	if dns.rcode == dnsRcodeNameError && dns.recursion_available {
		return "", nil, &DNSError{Err: noSuchHost, Name: name}
	}
	if dns.rcode != dnsRcodeSuccess {
		// None of the error codes make sense
		// for the query we sent.  If we didn't get
		// a name error and we didn't get success,
		// the server is behaving incorrectly.
		return "", nil, &DNSError{Err: "server misbehaving", Name: name, Server: server}
	}

	// Look for the name.
	// Presotto says it's okay to assume that servers listed in
	// /etc/resolv.conf are recursive resolvers.
	// We asked for recursion, so it should have included
	// all the answers we need in this one packet.
Cname:
	for cnameloop := 0; cnameloop < 10; cnameloop++ {
		addrs = addrs[0:0]
		for _, rr := range dns.answer {
			if _, justHeader := rr.(*dnsRR_Header); justHeader {
				// Corrupt record: we only have a
				// header. That header might say it's
				// of type qtype, but we don't
				// actually have it. Skip.
				continue
			}
			h := rr.Header()
			if h.Class == dnsClassINET && h.Name == name {
				switch h.Rrtype {
				case qtype:
					addrs = append(addrs, rr)
				case dnsTypeCNAME:
					// redirect to cname
					name = rr.(*dnsRR_CNAME).Cname
					continue Cname
				}
			}
		}
		if len(addrs) == 0 {
			return "", nil, &DNSError{Err: noSuchHost, Name: name, Server: server}
		}
		return name, addrs, nil
	}

	return "", nil, &DNSError{Err: "too many redirects", Name: name, Server: server}
}

func isDomainName(s string) bool {
	// See RFC 1035, RFC 3696.
	if len(s) == 0 {
		return false
	}
	if len(s) > 255 {
		return false
	}

	last := byte('.')
	ok := false // Ok once we've seen a letter.
	partlen := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		switch {
		default:
			return false
		case 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' || c == '_':
			ok = true
			partlen++
		case '0' <= c && c <= '9':
			// fine
			partlen++
		case c == '-':
			// Byte before dash cannot be dot.
			if last == '.' {
				return false
			}
			partlen++
		case c == '.':
			// Byte before dot cannot be dot, dash.
			if last == '.' || last == '-' {
				return false
			}
			if partlen > 63 || partlen == 0 {
				return false
			}
			partlen = 0
		}
		last = c
	}
	if last == '-' || partlen > 63 {
		return false
	}

	return ok
}

// An SRV represents a single DNS SRV record.
type SRV struct {
	Target   string
	Port     uint16
	Priority uint16
	Weight   uint16
}

// byPriorityWeight sorts SRV records by ascending priority and weight.
type byPriorityWeight []*SRV

func (s byPriorityWeight) Len() int { return len(s) }

func (s byPriorityWeight) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

func (s byPriorityWeight) Less(i, j int) bool {
	return s[i].Priority < s[j].Priority ||
		(s[i].Priority == s[j].Priority && s[i].Weight < s[j].Weight)
}

// shuffleByWeight shuffles SRV records by weight using the algorithm
// described in RFC 2782.
func (addrs byPriorityWeight) shuffleByWeight() {
	sum := 0
	for _, addr := range addrs {
		sum += int(addr.Weight)
	}
	for sum > 0 && len(addrs) > 1 {
		s := 0
		n := rand.Intn(sum)
		for i := range addrs {
			s += int(addrs[i].Weight)
			if s > n {
				if i > 0 {
					t := addrs[i]
					copy(addrs[1:i+1], addrs[0:i])
					addrs[0] = t
				}
				break
			}
		}
		sum -= int(addrs[0].Weight)
		addrs = addrs[1:]
	}
}

// sort reorders SRV records as specified in RFC 2782.
func (addrs byPriorityWeight) sort() {
	sort.Sort(addrs)
	i := 0
	for j := 1; j < len(addrs); j++ {
		if addrs[i].Priority != addrs[j].Priority {
			addrs[i:j].shuffleByWeight()
			i = j
		}
	}
	addrs[i:].shuffleByWeight()
}

// An MX represents a single DNS MX record.
type MX struct {
	Host string
	Pref uint16
}

// byPref implements sort.Interface to sort MX records by preference
type byPref []*MX

func (s byPref) Len() int { return len(s) }

func (s byPref) Less(i, j int) bool { return s[i].Pref < s[j].Pref }

func (s byPref) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// sort reorders MX records as specified in RFC 5321.
func (s byPref) sort() {
	for i := range s {
		j := rand.Intn(i + 1)
		s[i], s[j] = s[j], s[i]
	}
	sort.Sort(s)
}

// An NS represents a single DNS NS record.
type NS struct {
	Host string
}
<commit_msg>net: avoid array copy when shuffling SRV records<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 net

import (
	"math/rand"
	"sort"
)

// DNSError represents a DNS lookup error.
type DNSError struct {
	Err       string // description of the error
	Name      string // name looked for
	Server    string // server used
	IsTimeout bool
}

func (e *DNSError) Error() string {
	if e == nil {
		return "<nil>"
	}
	s := "lookup " + e.Name
	if e.Server != "" {
		s += " on " + e.Server
	}
	s += ": " + e.Err
	return s
}

func (e *DNSError) Timeout() bool   { return e.IsTimeout }
func (e *DNSError) Temporary() bool { return e.IsTimeout }

const noSuchHost = "no such host"

// reverseaddr returns the in-addr.arpa. or ip6.arpa. hostname of the IP
// address addr suitable for rDNS (PTR) record lookup or an error if it fails
// to parse the IP address.
func reverseaddr(addr string) (arpa string, err error) {
	ip := ParseIP(addr)
	if ip == nil {
		return "", &DNSError{Err: "unrecognized address", Name: addr}
	}
	if ip.To4() != nil {
		return itoa(int(ip[15])) + "." + itoa(int(ip[14])) + "." + itoa(int(ip[13])) + "." +
			itoa(int(ip[12])) + ".in-addr.arpa.", nil
	}
	// Must be IPv6
	buf := make([]byte, 0, len(ip)*4+len("ip6.arpa."))
	// Add it, in reverse, to the buffer
	for i := len(ip) - 1; i >= 0; i-- {
		v := ip[i]
		buf = append(buf, hexDigit[v&0xF])
		buf = append(buf, '.')
		buf = append(buf, hexDigit[v>>4])
		buf = append(buf, '.')
	}
	// Append "ip6.arpa." and return (buf already has the final .)
	buf = append(buf, "ip6.arpa."...)
	return string(buf), nil
}

// Find answer for name in dns message.
// On return, if err == nil, addrs != nil.
func answer(name, server string, dns *dnsMsg, qtype uint16) (cname string, addrs []dnsRR, err error) {
	addrs = make([]dnsRR, 0, len(dns.answer))

	if dns.rcode == dnsRcodeNameError && dns.recursion_available {
		return "", nil, &DNSError{Err: noSuchHost, Name: name}
	}
	if dns.rcode != dnsRcodeSuccess {
		// None of the error codes make sense
		// for the query we sent.  If we didn't get
		// a name error and we didn't get success,
		// the server is behaving incorrectly.
		return "", nil, &DNSError{Err: "server misbehaving", Name: name, Server: server}
	}

	// Look for the name.
	// Presotto says it's okay to assume that servers listed in
	// /etc/resolv.conf are recursive resolvers.
	// We asked for recursion, so it should have included
	// all the answers we need in this one packet.
Cname:
	for cnameloop := 0; cnameloop < 10; cnameloop++ {
		addrs = addrs[0:0]
		for _, rr := range dns.answer {
			if _, justHeader := rr.(*dnsRR_Header); justHeader {
				// Corrupt record: we only have a
				// header. That header might say it's
				// of type qtype, but we don't
				// actually have it. Skip.
				continue
			}
			h := rr.Header()
			if h.Class == dnsClassINET && h.Name == name {
				switch h.Rrtype {
				case qtype:
					addrs = append(addrs, rr)
				case dnsTypeCNAME:
					// redirect to cname
					name = rr.(*dnsRR_CNAME).Cname
					continue Cname
				}
			}
		}
		if len(addrs) == 0 {
			return "", nil, &DNSError{Err: noSuchHost, Name: name, Server: server}
		}
		return name, addrs, nil
	}

	return "", nil, &DNSError{Err: "too many redirects", Name: name, Server: server}
}

func isDomainName(s string) bool {
	// See RFC 1035, RFC 3696.
	if len(s) == 0 {
		return false
	}
	if len(s) > 255 {
		return false
	}

	last := byte('.')
	ok := false // Ok once we've seen a letter.
	partlen := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		switch {
		default:
			return false
		case 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' || c == '_':
			ok = true
			partlen++
		case '0' <= c && c <= '9':
			// fine
			partlen++
		case c == '-':
			// Byte before dash cannot be dot.
			if last == '.' {
				return false
			}
			partlen++
		case c == '.':
			// Byte before dot cannot be dot, dash.
			if last == '.' || last == '-' {
				return false
			}
			if partlen > 63 || partlen == 0 {
				return false
			}
			partlen = 0
		}
		last = c
	}
	if last == '-' || partlen > 63 {
		return false
	}

	return ok
}

// An SRV represents a single DNS SRV record.
type SRV struct {
	Target   string
	Port     uint16
	Priority uint16
	Weight   uint16
}

// byPriorityWeight sorts SRV records by ascending priority and weight.
type byPriorityWeight []*SRV

func (s byPriorityWeight) Len() int { return len(s) }

func (s byPriorityWeight) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

func (s byPriorityWeight) Less(i, j int) bool {
	return s[i].Priority < s[j].Priority ||
		(s[i].Priority == s[j].Priority && s[i].Weight < s[j].Weight)
}

// shuffleByWeight shuffles SRV records by weight using the algorithm
// described in RFC 2782.
func (addrs byPriorityWeight) shuffleByWeight() {
	sum := 0
	for _, addr := range addrs {
		sum += int(addr.Weight)
	}
	for sum > 0 && len(addrs) > 1 {
		s := 0
		n := rand.Intn(sum)
		for i := range addrs {
			s += int(addrs[i].Weight)
			if s > n {
				if i > 0 {
					addrs[0], addrs[i] = addrs[i], addrs[0]
				}
				break
			}
		}
		sum -= int(addrs[0].Weight)
		addrs = addrs[1:]
	}
}

// sort reorders SRV records as specified in RFC 2782.
func (addrs byPriorityWeight) sort() {
	sort.Sort(addrs)
	i := 0
	for j := 1; j < len(addrs); j++ {
		if addrs[i].Priority != addrs[j].Priority {
			addrs[i:j].shuffleByWeight()
			i = j
		}
	}
	addrs[i:].shuffleByWeight()
}

// An MX represents a single DNS MX record.
type MX struct {
	Host string
	Pref uint16
}

// byPref implements sort.Interface to sort MX records by preference
type byPref []*MX

func (s byPref) Len() int { return len(s) }

func (s byPref) Less(i, j int) bool { return s[i].Pref < s[j].Pref }

func (s byPref) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// sort reorders MX records as specified in RFC 5321.
func (s byPref) sort() {
	for i := range s {
		j := rand.Intn(i + 1)
		s[i], s[j] = s[j], s[i]
	}
	sort.Sort(s)
}

// An NS represents a single DNS NS record.
type NS struct {
	Host string
}
<|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 rpc

/*
	Some HTML presented at http://machine:port/debug/rpc
	Lists services, their methods, and some statistics, still rudimentary.
*/

import (
	"fmt"
	"net/http"
	"sort"
	"text/template"
)

const debugText = `

	<title>Services</title>
	{{range .}}

	Service {{.Name}}

		Method	Calls
		{{range .Method}}

			{{.Name}}({{.Type.ArgType}}, {{.Type.ReplyType}}) error
			{{.Type.NumCalls}}

		{{end}}

	{{end}}

	`

var debug = template.Must(template.New("RPC debug").Parse(debugText))

// If set, print log statements for internal and I/O errors.
var debugLog = false

type debugMethod struct {
	Type *methodType
	Name string
}

type methodArray []debugMethod

type debugService struct {
	Service *service
	Name    string
	Method  methodArray
}

type serviceArray []debugService

func (s serviceArray) Len() int           { return len(s) }
func (s serviceArray) Less(i, j int) bool { return s[i].Name < s[j].Name }
func (s serviceArray) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

func (m methodArray) Len() int           { return len(m) }
func (m methodArray) Less(i, j int) bool { return m[i].Name < m[j].Name }
func (m methodArray) Swap(i, j int)      { m[i], m[j] = m[j], m[i] }

type debugHTTP struct {
	*Server
}

// Runs at /debug/rpc
func (server debugHTTP) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// Build a sorted version of the data.
	var services = make(serviceArray, len(server.serviceMap))
	i := 0
	server.mu.Lock()
	for sname, service := range server.serviceMap {
		services[i] = debugService{service, sname, make(methodArray, len(service.method))}
		j := 0
		for mname, method := range service.method {
			services[i].Method[j] = debugMethod{method, mname}
			j++
		}
		sort.Sort(services[i].Method)
		i++
	}
	server.mu.Unlock()
	sort.Sort(services)
	err := debug.Execute(w, services)
	if err != nil {
		fmt.Fprintln(w, "rpc: error executing template:", err.Error())
	}
}
<commit_msg>net/rpc: use html/template to render html<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 rpc

/*
	Some HTML presented at http://machine:port/debug/rpc
	Lists services, their methods, and some statistics, still rudimentary.
*/

import (
	"fmt"
	"html/template"
	"net/http"
	"sort"
)

const debugText = `

	<title>Services</title>
	{{range .}}

	Service {{.Name}}

		Method	Calls
		{{range .Method}}

			{{.Name}}({{.Type.ArgType}}, {{.Type.ReplyType}}) error
			{{.Type.NumCalls}}

		{{end}}

	{{end}}

	`

var debug = template.Must(template.New("RPC debug").Parse(debugText))

// If set, print log statements for internal and I/O errors.
var debugLog = false

type debugMethod struct {
	Type *methodType
	Name string
}

type methodArray []debugMethod

type debugService struct {
	Service *service
	Name    string
	Method  methodArray
}

type serviceArray []debugService

func (s serviceArray) Len() int           { return len(s) }
func (s serviceArray) Less(i, j int) bool { return s[i].Name < s[j].Name }
func (s serviceArray) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

func (m methodArray) Len() int           { return len(m) }
func (m methodArray) Less(i, j int) bool { return m[i].Name < m[j].Name }
func (m methodArray) Swap(i, j int)      { m[i], m[j] = m[j], m[i] }

type debugHTTP struct {
	*Server
}

// Runs at /debug/rpc
func (server debugHTTP) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// Build a sorted version of the data.
	var services = make(serviceArray, len(server.serviceMap))
	i := 0
	server.mu.Lock()
	for sname, service := range server.serviceMap {
		services[i] = debugService{service, sname, make(methodArray, len(service.method))}
		j := 0
		for mname, method := range service.method {
			services[i].Method[j] = debugMethod{method, mname}
			j++
		}
		sort.Sort(services[i].Method)
		i++
	}
	server.mu.Unlock()
	sort.Sort(services)
	err := debug.Execute(w, services)
	if err != nil {
		fmt.Fprintln(w, "rpc: error executing template:", err.Error())
	}
}
<|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.

// +build darwin freebsd linux netbsd openbsd windows

package os

import (
	"syscall"
)

// The only signal values guaranteed to be present on all systems
// are Interrupt (send the process an interrupt) and Kill (force
// the process to exit).
var (
	Interrupt Signal = syscall.SIGINT
	Kill      Signal = syscall.SIGKILL
)

func startProcess(name string, argv []string, attr *ProcAttr) (p *Process, err error) {
	// If there is no SysProcAttr (ie. no Chroot or changed
	// UID/GID), double-check existence of the directory we want
	// to chdir into.  We can make the error clearer this way.
	if attr != nil && attr.Sys == nil && attr.Dir != "" {
		if _, err := Stat(attr.Dir); err != nil {
			pe := err.(*PathError)
			pe.Op = "chdir"
			return nil, pe
		}
	}

	sysattr := &syscall.ProcAttr{
		Dir: attr.Dir,
		Env: attr.Env,
		Sys: attr.Sys,
	}
	if sysattr.Env == nil {
		sysattr.Env = Environ()
	}
	for _, f := range attr.Files {
		sysattr.Files = append(sysattr.Files, f.Fd())
	}

	pid, h, e := syscall.StartProcess(name, argv, sysattr)
	if e != nil {
		return nil, &PathError{"fork/exec", name, e}
	}
	return newProcess(pid, h), nil
}

func (p *Process) kill() error {
	return p.Signal(Kill)
}

// ProcessState stores information about a process, as reported by Wait.
type ProcessState struct {
	pid    int                // The process's id.
	status syscall.WaitStatus // System-dependent status info.
	rusage *syscall.Rusage
}

// Pid returns the process id of the exited process.
func (p *ProcessState) Pid() int {
	return p.pid
}

func (p *ProcessState) exited() bool {
	return p.status.Exited()
}

func (p *ProcessState) success() bool {
	return p.status.ExitStatus() == 0
}

func (p *ProcessState) sys() interface{} {
	return p.status
}

func (p *ProcessState) sysUsage() interface{} {
	return p.rusage
}

// Convert i to decimal string.
func itod(i int) string {
	if i == 0 {
		return "0"
	}

	u := uint64(i)
	if i < 0 {
		u = -u
	}

	// Assemble decimal in reverse order.
	var b [32]byte
	bp := len(b)
	for ; u > 0; u /= 10 {
		bp--
		b[bp] = byte(u%10) + '0'
	}

	if i < 0 {
		bp--
		b[bp] = '-'
	}

	return string(b[bp:])
}

func (p *ProcessState) String() string {
	if p == nil {
		return "<nil>"
	}
	status := p.Sys().(syscall.WaitStatus)
	res := ""
	switch {
	case status.Exited():
		res = "exit status " + itod(status.ExitStatus())
	case status.Signaled():
		res = "signal " + itod(int(status.Signal()))
	case status.Stopped():
		res = "stop signal " + itod(int(status.StopSignal()))
		if status.StopSignal() == syscall.SIGTRAP && status.TrapCause() != 0 {
			res += " (trap " + itod(status.TrapCause()) + ")"
		}
	case status.Continued():
		res = "continued"
	}
	if status.CoreDump() {
		res += " (core dumped)"
	}
	return res
}
<commit_msg>os: use signal strings where possible in ProcessState.String<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.

// +build darwin freebsd linux netbsd openbsd windows

package os

import (
	"syscall"
)

// The only signal values guaranteed to be present on all systems
// are Interrupt (send the process an interrupt) and Kill (force
// the process to exit).
var (
	Interrupt Signal = syscall.SIGINT
	Kill      Signal = syscall.SIGKILL
)

func startProcess(name string, argv []string, attr *ProcAttr) (p *Process, err error) {
	// If there is no SysProcAttr (ie. no Chroot or changed
	// UID/GID), double-check existence of the directory we want
	// to chdir into.  We can make the error clearer this way.
	if attr != nil && attr.Sys == nil && attr.Dir != "" {
		if _, err := Stat(attr.Dir); err != nil {
			pe := err.(*PathError)
			pe.Op = "chdir"
			return nil, pe
		}
	}

	sysattr := &syscall.ProcAttr{
		Dir: attr.Dir,
		Env: attr.Env,
		Sys: attr.Sys,
	}
	if sysattr.Env == nil {
		sysattr.Env = Environ()
	}
	for _, f := range attr.Files {
		sysattr.Files = append(sysattr.Files, f.Fd())
	}

	pid, h, e := syscall.StartProcess(name, argv, sysattr)
	if e != nil {
		return nil, &PathError{"fork/exec", name, e}
	}
	return newProcess(pid, h), nil
}

func (p *Process) kill() error {
	return p.Signal(Kill)
}

// ProcessState stores information about a process, as reported by Wait.
type ProcessState struct {
	pid    int                // The process's id.
	status syscall.WaitStatus // System-dependent status info.
	rusage *syscall.Rusage
}

// Pid returns the process id of the exited process.
func (p *ProcessState) Pid() int {
	return p.pid
}

func (p *ProcessState) exited() bool {
	return p.status.Exited()
}

func (p *ProcessState) success() bool {
	return p.status.ExitStatus() == 0
}

func (p *ProcessState) sys() interface{} {
	return p.status
}

func (p *ProcessState) sysUsage() interface{} {
	return p.rusage
}

// Convert i to decimal string.
func itod(i int) string {
	if i == 0 {
		return "0"
	}

	u := uint64(i)
	if i < 0 {
		u = -u
	}

	// Assemble decimal in reverse order.
	var b [32]byte
	bp := len(b)
	for ; u > 0; u /= 10 {
		bp--
		b[bp] = byte(u%10) + '0'
	}

	if i < 0 {
		bp--
		b[bp] = '-'
	}

	return string(b[bp:])
}

func (p *ProcessState) String() string {
	if p == nil {
		return "<nil>"
	}
	status := p.Sys().(syscall.WaitStatus)
	res := ""
	switch {
	case status.Exited():
		res = "exit status " + itod(status.ExitStatus())
	case status.Signaled():
		res = "signal: " + status.Signal().String()
	case status.Stopped():
		res = "stop signal: " + status.StopSignal().String()
		if status.StopSignal() == syscall.SIGTRAP && status.TrapCause() != 0 {
			res += " (trap " + itod(status.TrapCause()) + ")"
		}
	case status.Continued():
		res = "continued"
	}
	if status.CoreDump() {
		res += " (core dumped)"
	}
	return res
}
<|endoftext|>


<commit_before>// Copyright (C) 2019 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

package protocol

import (
	"bytes"
	"fmt"
	"reflect"
	"regexp"
	"strings"
	"sync"
	"testing"

	"github.com/syncthing/syncthing/lib/rand"
	"github.com/syncthing/syncthing/lib/sha256"
)

func TestEnDecryptName(t *testing.T) {
	pattern := regexp.MustCompile(
		fmt.Sprintf("^[0-9A-V]%s/[0-9A-V]{2}/([0-9A-V]{%d}/)*[0-9A-V]{1,%d}$",
			regexp.QuoteMeta(encryptedDirExtension),
			maxPathComponent, maxPathComponent-1))

	makeName := func(n int) string {
		b := make([]byte, n)
		for i := range b {
			b[i] = byte('a' + i%26)
		}
		return string(b)
	}

	var key [32]byte
	cases := []string{
		"",
		"foo",
		"a longer name/with/slashes and spaces",
		makeName(maxPathComponent),
		makeName(1 + maxPathComponent),
		makeName(2 * maxPathComponent),
		makeName(1 + 2*maxPathComponent),
	}
	for _, tc := range cases {
		var prev string
		for i := 0; i < 5; i++ {
			enc := encryptName(tc, &key)
			if prev != "" && prev != enc {
				t.Error("name should always encrypt the same")
			}
			prev = enc
			if tc != "" && strings.Contains(enc, tc) {
				t.Error("shouldn't contain plaintext")
			}
			if !pattern.MatchString(enc) {
				t.Fatalf("encrypted name %s doesn't match %s",
					enc, pattern)
			}

			dec, err := decryptName(enc, &key)
			if err != nil {
				t.Error(err)
			}
			if dec != tc {
				t.Error("mismatch after decryption")
			}
			t.Logf("%q encrypts as %q", tc, enc)
		}
	}
}

func TestDecryptNameInvalid(t *testing.T) {
	key := new([32]byte)
	for _, c := range []string{
		"T.syncthing-enc/OD",
		"T.syncthing-enc/OD/",
		"T.wrong-extension/OD/PHVDD67S7FI2K5QQMPSOFSK",
		"OD/PHVDD67S7FI2K5QQMPSOFSK",
	} {
		if _, err := decryptName(c, key); err == nil {
			t.Errorf("no error for %q", c)
		}
	}
}

func TestEnDecryptBytes(t *testing.T) {
	var key [32]byte
	cases := [][]byte{
		{},
		{1, 2, 3, 4, 5},
	}
	for _, tc := range cases {
		var prev []byte
		for i := 0; i < 5; i++ {
			enc := encryptBytes(tc, &key)
			if bytes.Equal(enc, prev) {
				t.Error("encryption should not repeat")
			}
			prev = enc
			if len(tc) > 0 && bytes.Contains(enc, tc) {
				t.Error("shouldn't contain plaintext")
			}
			dec, err := DecryptBytes(enc, &key)
			if err != nil {
				t.Error(err)
			}
			if !bytes.Equal(dec, tc) {
				t.Error("mismatch after decryption")
			}
		}
	}
}

func encFileInfo() FileInfo {
	return FileInfo{
		Name:        "hello",
		Size:        45,
		Permissions: 0755,
		ModifiedS:   8080,
		Blocks: []BlockInfo{
			{
				Offset: 0,
				Size:   45,
				Hash:   []byte{1, 2, 3},
			},
			{
				Offset: 45,
				Size:   45,
				Hash:   []byte{1, 2, 3},
			},
		},
	}
}

func TestEnDecryptFileInfo(t *testing.T) {
	var key [32]byte
	fi := encFileInfo()

	enc := encryptFileInfo(fi, &key)
	if bytes.Equal(enc.Blocks[0].Hash, enc.Blocks[1].Hash) {
		t.Error("block hashes should not repeat when on different offsets")
	}
	if enc.RawBlockSize < MinBlockSize {
		t.Error("Too small raw block size:", enc.RawBlockSize)
	}
	again := encryptFileInfo(fi, &key)
	if !bytes.Equal(enc.Blocks[0].Hash, again.Blocks[0].Hash) {
		t.Error("block hashes should remain stable (0)")
	}
	if !bytes.Equal(enc.Blocks[1].Hash, again.Blocks[1].Hash) {
		t.Error("block hashes should remain stable (1)")
	}

	dec, err := DecryptFileInfo(enc, &key)
	if err != nil {
		t.Error(err)
	}
	if !reflect.DeepEqual(fi, dec) {
		t.Error("mismatch after decryption")
	}
}

func TestEncryptedFileInfoConsistency(t *testing.T) {
	var key [32]byte
	files := []FileInfo{
		encFileInfo(),
		encFileInfo(),
	}
	files[1].SetIgnored()
	for i, f := range files {
		enc := encryptFileInfo(f, &key)
		if err := checkFileInfoConsistency(enc); err != nil {
			t.Errorf("%v: %v", i, err)
		}
	}
}

func TestIsEncryptedParent(t *testing.T) {
	comp := rand.String(maxPathComponent)
	cases := []struct {
		path string
		is   bool
	}{
		{"", false},
		{".", false},
		{"/", false},
		{"12" + encryptedDirExtension, false},
		{"1" + encryptedDirExtension, true},
		{"1" + encryptedDirExtension + "/b", false},
		{"1" + encryptedDirExtension + "/bc", true},
		{"1" + encryptedDirExtension + "/bcd", false},
		{"1" + encryptedDirExtension + "/bc/foo", false},
		{"1.12/22", false},
		{"1" + encryptedDirExtension + "/bc/" + comp, true},
		{"1" + encryptedDirExtension + "/bc/" + comp + "/" + comp, true},
		{"1" + encryptedDirExtension + "/bc/" + comp + "a", false},
		{"1" + encryptedDirExtension + "/bc/" + comp + "/a/" + comp, false},
	}
	for _, tc := range cases {
		if res := IsEncryptedParent(strings.Split(tc.path, "/")); res != tc.is {
			t.Errorf("%v: got %v, expected %v", tc.path, res, tc.is)
		}
	}
}

var benchmarkFileKey struct {
	key [keySize]byte
	sync.Once
}

func BenchmarkFileKey(b *testing.B) {
	benchmarkFileKey.Do(func() {
		sha256.SelectAlgo()
		rand.Read(benchmarkFileKey.key[:])
	})

	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		FileKey("a_kind_of_long_filename.ext", &benchmarkFileKey.key)
	}
}
<commit_msg>lib/protocol: Test to lock down encryption key derivation<commit_after>// Copyright (C) 2019 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

package protocol

import (
	"bytes"
	"fmt"
	"reflect"
	"regexp"
	"strings"
	"sync"
	"testing"

	"github.com/syncthing/syncthing/lib/rand"
	"github.com/syncthing/syncthing/lib/sha256"
)

func TestEnDecryptName(t *testing.T) {
	pattern := regexp.MustCompile(
		fmt.Sprintf("^[0-9A-V]%s/[0-9A-V]{2}/([0-9A-V]{%d}/)*[0-9A-V]{1,%d}$",
			regexp.QuoteMeta(encryptedDirExtension),
			maxPathComponent, maxPathComponent-1))

	makeName := func(n int) string {
		b := make([]byte, n)
		for i := range b {
			b[i] = byte('a' + i%26)
		}
		return string(b)
	}

	var key [32]byte
	cases := []string{
		"",
		"foo",
		"a longer name/with/slashes and spaces",
		makeName(maxPathComponent),
		makeName(1 + maxPathComponent),
		makeName(2 * maxPathComponent),
		makeName(1 + 2*maxPathComponent),
	}
	for _, tc := range cases {
		var prev string
		for i := 0; i < 5; i++ {
			enc := encryptName(tc, &key)
			if prev != "" && prev != enc {
				t.Error("name should always encrypt the same")
			}
			prev = enc
			if tc != "" && strings.Contains(enc, tc) {
				t.Error("shouldn't contain plaintext")
			}
			if !pattern.MatchString(enc) {
				t.Fatalf("encrypted name %s doesn't match %s",
					enc, pattern)
			}

			dec, err := decryptName(enc, &key)
			if err != nil {
				t.Error(err)
			}
			if dec != tc {
				t.Error("mismatch after decryption")
			}
			t.Logf("%q encrypts as %q", tc, enc)
		}
	}
}

func TestKeyDerivation(t *testing.T) {
	folderKey := KeyFromPassword("my folder", "my password")
	encryptedName := encryptDeterministic([]byte("filename.txt"), folderKey, nil)
	if base32Hex.EncodeToString(encryptedName) != "3T5957I4IOA20VEIEER6JSQG0PEPIRV862II3K7LOF75Q" {
		t.Error("encrypted name mismatch")
	}

	fileKey := FileKey("filename.txt", folderKey)
	// fmt.Println(base32Hex.EncodeToString(encryptBytes([]byte("hello world"), fileKey))) => A1IPD...
	const encrypted = `A1IPD28ISL7VNPRSSSQM2L31L3IJPC08283RO89J5UG0TI9P38DO9RFGK12DK0KD7PKQP6U51UL2B6H96O`
	bs, _ := base32Hex.DecodeString(encrypted)
	dec, err := DecryptBytes(bs, fileKey)
	if err != nil {
		t.Error(err)
	}
	if string(dec) != "hello world" {
		t.Error("decryption mismatch")
	}
}

func TestDecryptNameInvalid(t *testing.T) {
	key := new([32]byte)
	for _, c := range []string{
		"T.syncthing-enc/OD",
		"T.syncthing-enc/OD/",
		"T.wrong-extension/OD/PHVDD67S7FI2K5QQMPSOFSK",
		"OD/PHVDD67S7FI2K5QQMPSOFSK",
	} {
		if _, err := decryptName(c, key); err == nil {
			t.Errorf("no error for %q", c)
		}
	}
}

func TestEnDecryptBytes(t *testing.T) {
	var key [32]byte
	cases := [][]byte{
		{},
		{1, 2, 3, 4, 5},
	}
	for _, tc := range cases {
		var prev []byte
		for i := 0; i < 5; i++ {
			enc := encryptBytes(tc, &key)
			if bytes.Equal(enc, prev) {
				t.Error("encryption should not repeat")
			}
			prev = enc
			if len(tc) > 0 && bytes.Contains(enc, tc) {
				t.Error("shouldn't contain plaintext")
			}
			dec, err := DecryptBytes(enc, &key)
			if err != nil {
				t.Error(err)
			}
			if !bytes.Equal(dec, tc) {
				t.Error("mismatch after decryption")
			}
		}
	}
}

func encFileInfo() FileInfo {
	return FileInfo{
		Name:        "hello",
		Size:        45,
		Permissions: 0755,
		ModifiedS:   8080,
		Blocks: []BlockInfo{
			{
				Offset: 0,
				Size:   45,
				Hash:   []byte{1, 2, 3},
			},
			{
				Offset: 45,
				Size:   45,
				Hash:   []byte{1, 2, 3},
			},
		},
	}
}

func TestEnDecryptFileInfo(t *testing.T) {
	var key [32]byte
	fi := encFileInfo()

	enc := encryptFileInfo(fi, &key)
	if bytes.Equal(enc.Blocks[0].Hash, enc.Blocks[1].Hash) {
		t.Error("block hashes should not repeat when on different offsets")
	}
	if enc.RawBlockSize < MinBlockSize {
		t.Error("Too small raw block size:", enc.RawBlockSize)
	}
	again := encryptFileInfo(fi, &key)
	if !bytes.Equal(enc.Blocks[0].Hash, again.Blocks[0].Hash) {
		t.Error("block hashes should remain stable (0)")
	}
	if !bytes.Equal(enc.Blocks[1].Hash, again.Blocks[1].Hash) {
		t.Error("block hashes should remain stable (1)")
	}

	dec, err := DecryptFileInfo(enc, &key)
	if err != nil {
		t.Error(err)
	}
	if !reflect.DeepEqual(fi, dec) {
		t.Error("mismatch after decryption")
	}
}

func TestEncryptedFileInfoConsistency(t *testing.T) {
	var key [32]byte
	files := []FileInfo{
		encFileInfo(),
		encFileInfo(),
	}
	files[1].SetIgnored()
	for i, f := range files {
		enc := encryptFileInfo(f, &key)
		if err := checkFileInfoConsistency(enc); err != nil {
			t.Errorf("%v: %v", i, err)
		}
	}
}

func TestIsEncryptedParent(t *testing.T) {
	comp := rand.String(maxPathComponent)
	cases := []struct {
		path string
		is   bool
	}{
		{"", false},
		{".", false},
		{"/", false},
		{"12" + encryptedDirExtension, false},
		{"1" + encryptedDirExtension, true},
		{"1" + encryptedDirExtension + "/b", false},
		{"1" + encryptedDirExtension + "/bc", true},
		{"1" + encryptedDirExtension + "/bcd", false},
		{"1" + encryptedDirExtension + "/bc/foo", false},
		{"1.12/22", false},
		{"1" + encryptedDirExtension + "/bc/" + comp, true},
		{"1" + encryptedDirExtension + "/bc/" + comp + "/" + comp, true},
		{"1" + encryptedDirExtension + "/bc/" + comp + "a", false},
		{"1" + encryptedDirExtension + "/bc/" + comp + "/a/" + comp, false},
	}
	for _, tc := range cases {
		if res := IsEncryptedParent(strings.Split(tc.path, "/")); res != tc.is {
			t.Errorf("%v: got %v, expected %v", tc.path, res, tc.is)
		}
	}
}

var benchmarkFileKey struct {
	key [keySize]byte
	sync.Once
}

func BenchmarkFileKey(b *testing.B) {
	benchmarkFileKey.Do(func() {
		sha256.SelectAlgo()
		rand.Read(benchmarkFileKey.key[:])
	})

	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		FileKey("a_kind_of_long_filename.ext", &benchmarkFileKey.key)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"fmt"
	"strings"
	"os"
	"strconv"
	"text/tabwriter"

	"github.com/nodetemple/cobra"
	"github.com/nodetemple/nodetemple/common"
)

const (
	cliName        = "nodectl"
	cliDescription = "nodectl is a command-line interface for an orchestration of CoreOS and Kubernetes cluster"
)

var (
	tabOut        *tabwriter.Writer
	globalFlags   = struct {
		Debug       bool
		Provider    string
	}{}

	cmdExitCode int

	cmdNodectl = &cobra.Command{
		Use:   fmt.Sprintf("%s [command]", cliName),
		Short: cliDescription,
	}
)

func init() {
	cmdNodectl.PersistentFlags().BoolVar(&globalFlags.Debug, "debug", envBool("debug", false), "Print out more debug information")
	cmdNodectl.PersistentFlags().StringVarP(&globalFlags.Provider, "provider", "p", envString("provider", common.DefaultProvider), "Provider to use when managing a cluster")

	tabOut = new(tabwriter.Writer)
	tabOut.Init(os.Stdout, 0, 8, 1, '\t', 0)

	cobra.EnablePrefixMatching = true
}

func main() {
	cmdNodectl.SetUsageFunc(usageFunc)
	cmdNodectl.SetUsageTemplate(`{{.UsageString}}`)

	cmdNodectl.SetHelpCommand(cmdHelp)
	cmdNodectl.SetHelpTemplate(`{{.UsageString}}`)

	cmdNodectl.Execute()
	os.Exit(cmdExitCode)
}

func stderr(format string, a ...interface{}) {
	out := fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stderr, strings.TrimSuffix(out, "\n"))
}

func stdout(format string, a ...interface{}) {
	out := fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stdout, strings.TrimSuffix(out, "\n"))
}

func runWrapper(cf func(cmd *cobra.Command, args []string) (exit int)) func(cmd *cobra.Command, args []string) {
	return func(cmd *cobra.Command, args []string) {
		cmdExitCode = cf(cmd, args)
	}
}

func envString(key, def string) string {
	envKey := strings.ToUpper(cliName + "_" + strings.Replace(key, "-", "_", -1))

	if env := os.Getenv(envKey); env != "" {
		return env
	}
	return def
}

func envBool(key string, def bool) bool {
	envKey := strings.ToUpper(cliName + "_" + strings.Replace(key, "-", "_", -1))

	if env := os.Getenv(envKey); env != "" {
		val, err := strconv.ParseBool(env)
		if err != nil {
			stderr("invalid value \"%v\" for \"%v\" (using default: %v): %v", env, key, def, err)
			return def
		}
		return val
	}
	return def
}

func envInt(key string, def int) int {
	envKey := strings.ToUpper(cliName + "_" + strings.Replace(key, "-", "_", -1))

	if env := os.Getenv(envKey); env != "" {
		val, err := strconv.Atoi(env)
		if err != nil {
			stderr("invalid value \"%v\" for \"%v\" (using default: %v): %v", env, key, def, err)
			return def
		}
		return val
	}
	return def
}
<commit_msg>Prefix stderr<commit_after>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"fmt"
	"strings"
	"os"
	"strconv"
	"text/tabwriter"

	"github.com/nodetemple/cobra"
	"github.com/nodetemple/nodetemple/common"
)

const (
	cliName        = "nodectl"
	cliDescription = "nodectl is a command-line interface for an orchestration of CoreOS and Kubernetes cluster"
)

var (
	tabOut        *tabwriter.Writer
	globalFlags   = struct {
		Debug       bool
		Provider    string
	}{}

	cmdExitCode int

	cmdNodectl = &cobra.Command{
		Use:   fmt.Sprintf("%s [command]", cliName),
		Short: cliDescription,
	}
)

func init() {
	cmdNodectl.PersistentFlags().BoolVar(&globalFlags.Debug, "debug", envBool("debug", false), "Print out more debug information")
	cmdNodectl.PersistentFlags().StringVarP(&globalFlags.Provider, "provider", "p", envString("provider", common.DefaultProvider), "Provider to use when managing a cluster")

	tabOut = new(tabwriter.Writer)
	tabOut.Init(os.Stdout, 0, 8, 1, '\t', 0)

	cobra.EnablePrefixMatching = true
}

func main() {
	cmdNodectl.SetUsageFunc(usageFunc)
	cmdNodectl.SetUsageTemplate(`{{.UsageString}}`)

	cmdNodectl.SetHelpCommand(cmdHelp)
	cmdNodectl.SetHelpTemplate(`{{.UsageString}}`)

	cmdNodectl.Execute()
	os.Exit(cmdExitCode)
}

func stderr(format string, a ...interface{}) {
	out := "Error:" + fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stderr, strings.TrimSuffix(out, "\n"))
}

func stdout(format string, a ...interface{}) {
	out := fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stdout, strings.TrimSuffix(out, "\n"))
}

func runWrapper(cf func(cmd *cobra.Command, args []string) (exit int)) func(cmd *cobra.Command, args []string) {
	return func(cmd *cobra.Command, args []string) {
		cmdExitCode = cf(cmd, args)
	}
}

func envString(key, def string) string {
	envKey := strings.ToUpper(cliName + "_" + strings.Replace(key, "-", "_", -1))

	if env := os.Getenv(envKey); env != "" {
		return env
	}
	return def
}

func envBool(key string, def bool) bool {
	envKey := strings.ToUpper(cliName + "_" + strings.Replace(key, "-", "_", -1))

	if env := os.Getenv(envKey); env != "" {
		val, err := strconv.ParseBool(env)
		if err != nil {
			stderr("invalid value \"%v\" for \"%v\" (using default: %v): %v", env, key, def, err)
			return def
		}
		return val
	}
	return def
}

func envInt(key string, def int) int {
	envKey := strings.ToUpper(cliName + "_" + strings.Replace(key, "-", "_", -1))

	if env := os.Getenv(envKey); env != "" {
		val, err := strconv.Atoi(env)
		if err != nil {
			stderr("invalid value \"%v\" for \"%v\" (using default: %v): %v", env, key, def, err)
			return def
		}
		return val
	}
	return def
}
<|endoftext|>