<commit_before>// Copyright 2014 The LUCI Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package local

import (
	"archive/zip"
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"golang.org/x/net/context"

	"github.com/luci/luci-go/cipd/client/cipd/common"
	"github.com/luci/luci-go/common/clock"
	"github.com/luci/luci-go/common/logging"
)

// PackageInstance represents a binary CIPD package file (with manifest inside).
type PackageInstance interface {
	// Close shuts down the package and its data provider.
	Close() error

	// Pin identifies package name and concreted instance ID of this package file.
	Pin() common.Pin

	// Files returns a list of files to deploy with the package.
	Files() []File

	// DataReader returns reader that reads raw package data.
	DataReader() io.ReadSeeker
}

// OpenInstance verifies the package and prepares it for extraction.
//
// It checks package SHA1 hash (must match instanceID, if it's given) and
// prepares a package instance for extraction. If the call succeeds,
// PackageInstance takes ownership of io.ReadSeeker. If it also implements
// io.Closer, it will be closed when package.Close() is called. If an error is
// returned, io.ReadSeeker remains unowned and caller is responsible for closing
// it (if required).
func OpenInstance(ctx context.Context, r io.ReadSeeker, instanceID string) (PackageInstance, error) {
	out := &packageInstance{data: r}
	if err := out.open(instanceID); err != nil {
		return nil, err
	}
	return out, nil
}

// OpenInstanceFile opens a package instance file on disk.
func OpenInstanceFile(ctx context.Context, path string, instanceID string) (inst PackageInstance, err error) {
	file, err := os.Open(path)
	if err != nil {
		return
	}
	inst, err = OpenInstance(ctx, file, instanceID)
	if err != nil {
		file.Close()
	}
	return
}

// ExtractInstance extracts all files from a package instance into a destination.
func ExtractInstance(ctx context.Context, inst PackageInstance, dest Destination) error {
	if err := dest.Begin(ctx); err != nil {
		return err
	}

	// Do not leave garbage around in case of a panic.
	needToEnd := true
	defer func() {
		if needToEnd {
			dest.End(ctx, false)
		}
	}()

	files := inst.Files()

	extractManifestFile := func(f File) (err error) {
		manifest, err := readManifestFile(f)
		if err != nil {
			return err
		}
		manifest.Files = make([]FileInfo, 0, len(files))
		for _, file := range files {
			// Do not put info about service .cipdpkg files into the manifest,
			// otherwise it becomes recursive and "size" property of manifest file
			// itself is not correct.
			if strings.HasPrefix(file.Name(), packageServiceDir+"/") {
				continue
			}
			fi := FileInfo{
				Name:       file.Name(),
				Size:       file.Size(),
				Executable: file.Executable(),
			}
			if file.Symlink() {
				target, err := file.SymlinkTarget()
				if err != nil {
					return err
				}
				fi.Symlink = target
			}
			manifest.Files = append(manifest.Files, fi)
		}
		out, err := dest.CreateFile(ctx, f.Name(), false)
		if err != nil {
			return err
		}
		defer func() {
			if closeErr := out.Close(); err == nil {
				err = closeErr
			}
		}()
		return writeManifest(&manifest, out)
	}

	extractSymlinkFile := func(f File) error {
		target, err := f.SymlinkTarget()
		if err != nil {
			return err
		}
		return dest.CreateSymlink(ctx, f.Name(), target)
	}

	extractRegularFile := func(f File) (err error) {
		out, err := dest.CreateFile(ctx, f.Name(), f.Executable())
		if err != nil {
			return err
		}
		defer func() {
			if closeErr := out.Close(); err == nil {
				err = closeErr
			}
		}()
		in, err := f.Open()
		if err != nil {
			return err
		}
		defer in.Close()
		_, err = io.Copy(out, in)
		return err
	}

	// Use file sizes for progress report calculation.
	totalSize := uint64(0)
	extractedSize := uint64(0)
	for _, f := range files {
		if !f.Symlink() {
			totalSize += f.Size()
		}
	}

	logging.Infof(ctx, "cipd: about to extract %.1f Mb", float64(totalSize)/1024.0/1024.0)

	// reportProgress print extraction progress, throttling the reports rate to
	// one per 5 sec.
	prevProgress := 1000 // >100%
	var prevReportTs time.Time
	reportProgress := func(read, total uint64) {
		now := clock.Now(ctx)
		progress := int(float64(read) * 100 / float64(total))
		if progress < prevProgress || read == total || now.Sub(prevReportTs) > 5*time.Second {
			logging.Infof(ctx, "cipd: extracting - %d%%", progress)
			prevReportTs = now
			prevProgress = progress
		}
	}

	var err error
	for _, f := range files {
		if err = ctx.Err(); err != nil {
			break
		}
		if f.Name() == manifestName {
			err = extractManifestFile(f)
			extractedSize += f.Size()
		} else if f.Symlink() {
			err = extractSymlinkFile(f)
		} else {
			err = extractRegularFile(f)
			extractedSize += f.Size()
		}
		if err != nil {
			break
		}
		reportProgress(extractedSize, totalSize)
	}

	needToEnd = false
	if err == nil {
		err = dest.End(ctx, true)
	} else {
		// Ignore error in 'End' and return the original error.
		dest.End(ctx, false)
	}

	return err
}

////////////////////////////////////////////////////////////////////////////////
// PackageInstance implementation.

type packageInstance struct {
	data       io.ReadSeeker
	instanceID string
	zip        *zip.Reader
	files      []File
	manifest   Manifest
}

// open reads the package data, verifies SHA1 hash and reads manifest.
func (inst *packageInstance) open(instanceID string) error {
	// Calculate SHA1 of the data to verify it matches expected instanceID.
	if _, err := inst.data.Seek(0, os.SEEK_SET); err != nil {
		return err
	}
	hash := sha1.New()
	if _, err := io.Copy(hash, inst.data); err != nil {
		return err
	}

	dataSize, err := inst.data.Seek(0, os.SEEK_CUR)
	if err != nil {
		return err
	}

	calculatedSHA1 := hex.EncodeToString(hash.Sum(nil))
	if instanceID != "" && instanceID != calculatedSHA1 {
		return fmt.Errorf("package SHA1 hash mismatch")
	}
	inst.instanceID = calculatedSHA1

	// Zip reader needs an io.ReaderAt. Try to sniff it from our io.ReadSeeker
	// before falling back to a generic (potentially slower) implementation. This
	// works if inst.data is actually an os.File (which happens quite often).
	reader, ok := inst.data.(io.ReaderAt)
	if !ok {
		reader = &readerAt{r: inst.data}
	}

	// List files and package manifest.
	inst.zip, err = zip.NewReader(reader, dataSize)
	if err != nil {
		return err
	}
	inst.files = make([]File, len(inst.zip.File))
	for i, zf := range inst.zip.File {
		inst.files[i] = &fileInZip{z: zf}
		if inst.files[i].Name() == manifestName {
			inst.manifest, err = readManifestFile(inst.files[i])
			if err != nil {
				return err
			}
		}
	}

	// Generate version_file if needed.
	if inst.manifest.VersionFile != "" {
		vf, err := makeVersionFile(inst.manifest.VersionFile, VersionFile{
			PackageName: inst.manifest.PackageName,
			InstanceID:  inst.instanceID,
		})
		if err != nil {
			return err
		}
		inst.files = append(inst.files, vf)
	}

	return nil
}

func (inst *packageInstance) Close() (err error) {
	if inst.data != nil {
		if closer, ok := inst.data.(io.Closer); ok {
			err = closer.Close()
		}
		inst.data = nil
	}
	inst.instanceID = ""
	inst.zip = nil
	inst.files = []File{}
	inst.manifest = Manifest{}
	return
}

func (inst *packageInstance) Pin() common.Pin {
	return common.Pin{
		PackageName: inst.manifest.PackageName,
		InstanceID:  inst.instanceID,
	}
}

func (inst *packageInstance) Files() []File             { return inst.files }
func (inst *packageInstance) DataReader() io.ReadSeeker { return inst.data }

////////////////////////////////////////////////////////////////////////////////
// Utilities.

// readManifestFile decodes manifest file zipped inside the package.
func readManifestFile(f File) (Manifest, error) {
	r, err := f.Open()
	if err != nil {
		return Manifest{}, err
	}
	defer r.Close()
	return readManifest(r)
}

// makeVersionFile returns File representing a JSON blob with info about package
// version. It's what's deployed at path specified in 'version_file' stanza in
// package definition YAML.
func makeVersionFile(relPath string, versionFile VersionFile) (File, error) {
	if !isCleanSlashPath(relPath) {
		return nil, fmt.Errorf("invalid version_file: %s", relPath)
	}
	blob, err := json.MarshalIndent(versionFile, "", "  ")
	if err != nil {
		return nil, err
	}
	return &blobFile{
		name: relPath,
		blob: blob,
	}, nil
}

// blobFile implements File on top of byte array with file data.
type blobFile struct {
	name string
	blob []byte
}

func (b *blobFile) Name() string                   { return b.name }
func (b *blobFile) Size() uint64                   { return uint64(len(b.blob)) }
func (b *blobFile) Executable() bool               { return false }
func (b *blobFile) Symlink() bool                  { return false }
func (b *blobFile) SymlinkTarget() (string, error) { return "", nil }

func (b *blobFile) Open() (io.ReadCloser, error) {
	return ioutil.NopCloser(bytes.NewReader(b.blob)), nil
}

////////////////////////////////////////////////////////////////////////////////
// File interface implementation via zip.File.

type fileInZip struct {
	z *zip.File
}

func (f *fileInZip) Name() string  { return f.z.Name }
func (f *fileInZip) Symlink() bool { return (f.z.Mode() & os.ModeSymlink) != 0 }

func (f *fileInZip) Executable() bool {
	if f.Symlink() {
		return false
	}
	return (f.z.Mode() & 0100) != 0
}

func (f *fileInZip) Size() uint64 {
	if f.Symlink() {
		return 0
	}
	return f.z.UncompressedSize64
}

func (f *fileInZip) SymlinkTarget() (string, error) {
	if !f.Symlink() {
		return "", fmt.Errorf("not a symlink: %s", f.Name())
	}
	r, err := f.z.Open()
	if err != nil {
		return "", err
	}
	defer r.Close()
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func (f *fileInZip) Open() (io.ReadCloser, error) {
	if f.Symlink() {
		return nil, fmt.Errorf("opening a symlink is not allowed: %s", f.Name())
	}
	return f.z.Open()
}

////////////////////////////////////////////////////////////////////////////////
// ReaderAt implementation via ReadSeeker. Not concurrency safe, moves file
// pointer around without any locking. Works OK in the context of OpenInstance
// function though (where OpenInstance takes sole ownership of io.ReadSeeker).

type readerAt struct {
	r io.ReadSeeker
}

func (r *readerAt) ReadAt(data []byte, off int64) (int, error) {
	_, err := r.r.Seek(off, os.SEEK_SET)
	if err != nil {
		return 0, err
	}
	return r.r.Read(data)
}
<commit_msg>cipd: Extract progress reported from ExtractInstance.<commit_after>// Copyright 2014 The LUCI Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package local

import (
	"archive/zip"
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/luci/luci-go/cipd/client/cipd/common"
	"github.com/luci/luci-go/common/clock"
	"github.com/luci/luci-go/common/logging"
)

// PackageInstance represents a binary CIPD package file (with manifest inside).
type PackageInstance interface {
	// Close shuts down the package and its data provider.
	Close() error

	// Pin identifies package name and concreted instance ID of this package file.
	Pin() common.Pin

	// Files returns a list of files to deploy with the package.
	Files() []File

	// DataReader returns reader that reads raw package data.
	DataReader() io.ReadSeeker
}

// OpenInstance verifies the package and prepares it for extraction.
//
// It checks package SHA1 hash (must match instanceID, if it's given) and
// prepares a package instance for extraction. If the call succeeds,
// PackageInstance takes ownership of io.ReadSeeker. If it also implements
// io.Closer, it will be closed when package.Close() is called. If an error is
// returned, io.ReadSeeker remains unowned and caller is responsible for closing
// it (if required).
func OpenInstance(ctx context.Context, r io.ReadSeeker, instanceID string) (PackageInstance, error) {
	out := &packageInstance{data: r}
	if err := out.open(instanceID); err != nil {
		return nil, err
	}
	return out, nil
}

// OpenInstanceFile opens a package instance file on disk.
func OpenInstanceFile(ctx context.Context, path string, instanceID string) (inst PackageInstance, err error) {
	file, err := os.Open(path)
	if err != nil {
		return
	}
	inst, err = OpenInstance(ctx, file, instanceID)
	if err != nil {
		file.Close()
	}
	return
}

// ExtractInstance extracts all files from a package instance into a destination.
func ExtractInstance(ctx context.Context, inst PackageInstance, dest Destination) error {
	if err := dest.Begin(ctx); err != nil {
		return err
	}

	// Do not leave garbage around in case of a panic.
	needToEnd := true
	defer func() {
		if needToEnd {
			dest.End(ctx, false)
		}
	}()

	files := inst.Files()
	progress := newProgressReporter(ctx, files)

	extractManifestFile := func(f File) (err error) {
		defer progress.advance(f)
		manifest, err := readManifestFile(f)
		if err != nil {
			return err
		}
		manifest.Files = make([]FileInfo, 0, len(files))
		for _, file := range files {
			// Do not put info about service .cipdpkg files into the manifest,
			// otherwise it becomes recursive and "size" property of manifest file
			// itself is not correct.
			if strings.HasPrefix(file.Name(), packageServiceDir+"/") {
				continue
			}
			fi := FileInfo{
				Name:       file.Name(),
				Size:       file.Size(),
				Executable: file.Executable(),
			}
			if file.Symlink() {
				target, err := file.SymlinkTarget()
				if err != nil {
					return err
				}
				fi.Symlink = target
			}
			manifest.Files = append(manifest.Files, fi)
		}
		out, err := dest.CreateFile(ctx, f.Name(), false)
		if err != nil {
			return err
		}
		defer func() {
			if closeErr := out.Close(); err == nil {
				err = closeErr
			}
		}()
		return writeManifest(&manifest, out)
	}

	extractSymlinkFile := func(f File) error {
		defer progress.advance(f)
		target, err := f.SymlinkTarget()
		if err != nil {
			return err
		}
		return dest.CreateSymlink(ctx, f.Name(), target)
	}

	extractRegularFile := func(f File) (err error) {
		defer progress.advance(f)
		out, err := dest.CreateFile(ctx, f.Name(), f.Executable())
		if err != nil {
			return err
		}
		defer func() {
			if closeErr := out.Close(); err == nil {
				err = closeErr
			}
		}()
		in, err := f.Open()
		if err != nil {
			return err
		}
		defer in.Close()
		_, err = io.Copy(out, in)
		return err
	}

	var err error
	for _, f := range files {
		if err = ctx.Err(); err != nil {
			break
		}
		switch {
		case f.Name() == manifestName:
			err = extractManifestFile(f)
		case f.Symlink():
			err = extractSymlinkFile(f)
		default:
			err = extractRegularFile(f)
		}
		if err != nil {
			break
		}
	}

	needToEnd = false
	if err == nil {
		err = dest.End(ctx, true)
	} else {
		// Ignore error in 'End' and return the original error.
		dest.End(ctx, false)
	}

	return err
}

// progressReporter periodically logs progress of the extraction.
//
// Can be shared by multiple goroutines.
type progressReporter struct {
	sync.Mutex

	ctx context.Context

	totalCount     uint64    // total number of files to extract
	totalSize      uint64    // total expected uncompressed size of files
	extractedCount uint64    // number of files extract so far
	extractedSize  uint64    // bytes uncompressed so far
	prevReport     time.Time // time when we did the last progress report
}

func newProgressReporter(ctx context.Context, files []File) *progressReporter {
	r := &progressReporter{ctx: ctx, totalCount: uint64(len(files))}
	for _, f := range files {
		if !f.Symlink() {
			r.totalSize += f.Size()
		}
	}
	if r.totalCount != 0 {
		logging.Infof(
			r.ctx, "cipd: about to extract %.1f Mb (%d files)",
			float64(r.totalSize)/1024.0/1024.0, r.totalCount)
	}
	return r
}

// advance moves the progress indicator, occasionally logging it.
func (r *progressReporter) advance(f File) {
	if r.totalCount == 0 {
		return
	}

	now := clock.Now(r.ctx)
	reportNow := false
	progress := 0

	// We don't count size of the symlinks toward total.
	var size uint64
	if !f.Symlink() {
		size = f.Size()
	}

	// Report progress on first and last 'advance' calls and each 2 sec.
	r.Lock()
	r.extractedSize += size
	r.extractedCount++
	if r.extractedCount == 1 || r.extractedCount == r.totalCount || now.Sub(r.prevReport) > 2*time.Second {
		reportNow = true
		if r.totalSize != 0 {
			progress = int(float64(r.extractedSize) * 100 / float64(r.totalSize))
		} else {
			progress = int(float64(r.extractedCount) * 100 / float64(r.totalCount))
		}
		r.prevReport = now
	}
	r.Unlock()

	if reportNow {
		logging.Infof(r.ctx, "cipd: extracting - %d%%", progress)
	}
}

////////////////////////////////////////////////////////////////////////////////
// PackageInstance implementation.

type packageInstance struct {
	data       io.ReadSeeker
	instanceID string
	zip        *zip.Reader
	files      []File
	manifest   Manifest
}

// open reads the package data, verifies SHA1 hash and reads manifest.
func (inst *packageInstance) open(instanceID string) error {
	// Calculate SHA1 of the data to verify it matches expected instanceID.
	if _, err := inst.data.Seek(0, os.SEEK_SET); err != nil {
		return err
	}
	hash := sha1.New()
	if _, err := io.Copy(hash, inst.data); err != nil {
		return err
	}

	dataSize, err := inst.data.Seek(0, os.SEEK_CUR)
	if err != nil {
		return err
	}

	calculatedSHA1 := hex.EncodeToString(hash.Sum(nil))
	if instanceID != "" && instanceID != calculatedSHA1 {
		return fmt.Errorf("package SHA1 hash mismatch")
	}
	inst.instanceID = calculatedSHA1

	// Zip reader needs an io.ReaderAt. Try to sniff it from our io.ReadSeeker
	// before falling back to a generic (potentially slower) implementation. This
	// works if inst.data is actually an os.File (which happens quite often).
	reader, ok := inst.data.(io.ReaderAt)
	if !ok {
		reader = &readerAt{r: inst.data}
	}

	// List files and package manifest.
	inst.zip, err = zip.NewReader(reader, dataSize)
	if err != nil {
		return err
	}
	inst.files = make([]File, len(inst.zip.File))
	for i, zf := range inst.zip.File {
		inst.files[i] = &fileInZip{z: zf}
		if inst.files[i].Name() == manifestName {
			inst.manifest, err = readManifestFile(inst.files[i])
			if err != nil {
				return err
			}
		}
	}

	// Generate version_file if needed.
	if inst.manifest.VersionFile != "" {
		vf, err := makeVersionFile(inst.manifest.VersionFile, VersionFile{
			PackageName: inst.manifest.PackageName,
			InstanceID:  inst.instanceID,
		})
		if err != nil {
			return err
		}
		inst.files = append(inst.files, vf)
	}

	return nil
}

func (inst *packageInstance) Close() (err error) {
	if inst.data != nil {
		if closer, ok := inst.data.(io.Closer); ok {
			err = closer.Close()
		}
		inst.data = nil
	}
	inst.instanceID = ""
	inst.zip = nil
	inst.files = []File{}
	inst.manifest = Manifest{}
	return
}

func (inst *packageInstance) Pin() common.Pin {
	return common.Pin{
		PackageName: inst.manifest.PackageName,
		InstanceID:  inst.instanceID,
	}
}

func (inst *packageInstance) Files() []File             { return inst.files }
func (inst *packageInstance) DataReader() io.ReadSeeker { return inst.data }

////////////////////////////////////////////////////////////////////////////////
// Utilities.

// readManifestFile decodes manifest file zipped inside the package.
func readManifestFile(f File) (Manifest, error) {
	r, err := f.Open()
	if err != nil {
		return Manifest{}, err
	}
	defer r.Close()
	return readManifest(r)
}

// makeVersionFile returns File representing a JSON blob with info about package
// version. It's what's deployed at path specified in 'version_file' stanza in
// package definition YAML.
func makeVersionFile(relPath string, versionFile VersionFile) (File, error) {
	if !isCleanSlashPath(relPath) {
		return nil, fmt.Errorf("invalid version_file: %s", relPath)
	}
	blob, err := json.MarshalIndent(versionFile, "", "  ")
	if err != nil {
		return nil, err
	}
	return &blobFile{
		name: relPath,
		blob: blob,
	}, nil
}

// blobFile implements File on top of byte array with file data.
type blobFile struct {
	name string
	blob []byte
}

func (b *blobFile) Name() string                   { return b.name }
func (b *blobFile) Size() uint64                   { return uint64(len(b.blob)) }
func (b *blobFile) Executable() bool               { return false }
func (b *blobFile) Symlink() bool                  { return false }
func (b *blobFile) SymlinkTarget() (string, error) { return "", nil }

func (b *blobFile) Open() (io.ReadCloser, error) {
	return ioutil.NopCloser(bytes.NewReader(b.blob)), nil
}

////////////////////////////////////////////////////////////////////////////////
// File interface implementation via zip.File.

type fileInZip struct {
	z *zip.File
}

func (f *fileInZip) Name() string  { return f.z.Name }
func (f *fileInZip) Symlink() bool { return (f.z.Mode() & os.ModeSymlink) != 0 }

func (f *fileInZip) Executable() bool {
	if f.Symlink() {
		return false
	}
	return (f.z.Mode() & 0100) != 0
}

func (f *fileInZip) Size() uint64 {
	if f.Symlink() {
		return 0
	}
	return f.z.UncompressedSize64
}

func (f *fileInZip) SymlinkTarget() (string, error) {
	if !f.Symlink() {
		return "", fmt.Errorf("not a symlink: %s", f.Name())
	}
	r, err := f.z.Open()
	if err != nil {
		return "", err
	}
	defer r.Close()
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func (f *fileInZip) Open() (io.ReadCloser, error) {
	if f.Symlink() {
		return nil, fmt.Errorf("opening a symlink is not allowed: %s", f.Name())
	}
	return f.z.Open()
}

////////////////////////////////////////////////////////////////////////////////
// ReaderAt implementation via ReadSeeker. Not concurrency safe, moves file
// pointer around without any locking. Works OK in the context of OpenInstance
// function though (where OpenInstance takes sole ownership of io.ReadSeeker).

type readerAt struct {
	r io.ReadSeeker
}

func (r *readerAt) ReadAt(data []byte, off int64) (int, error) {
	_, err := r.r.Seek(off, os.SEEK_SET)
	if err != nil {
		return 0, err
	}
	return r.r.Read(data)
}
<|endoftext|>


<commit_before>package iplist

import (
	"encoding/binary"
	"io"
	"net"
)

// The packed format is an 8 byte integer of the number of ranges. Then 20
// bytes per range, consisting of 4 byte packed IP being the lower bound IP of
// the range, then 4 bytes of the upper, inclusive bound, 8 bytes for the
// offset of the description from the end of the packed ranges, and 4 bytes
// for the length of the description. After these packed ranges, are the
// concatenated descriptions.

const (
	packedRangesOffset = 8
	packedRangeLen     = 20
)

func (me *IPList) WritePacked(w io.Writer) (err error) {
	descOffsets := make(map[string]int64, len(me.ranges))
	descs := make([]string, 0, len(me.ranges))
	var nextOffset int64
	// This is a little monadic, no?
	write := func(b []byte, expectedLen int) {
		if err != nil {
			return
		}
		var n int
		n, err = w.Write(b)
		if err != nil {
			return
		}
		if n != expectedLen {
			panic(n)
		}
	}
	var b [8]byte
	binary.LittleEndian.PutUint64(b[:], uint64(len(me.ranges)))
	write(b[:], 8)
	for _, r := range me.ranges {
		write(r.First.To4(), 4)
		write(r.Last.To4(), 4)
		descOff, ok := descOffsets[r.Description]
		if !ok {
			descOff = nextOffset
			descOffsets[r.Description] = descOff
			descs = append(descs, r.Description)
			nextOffset += int64(len(r.Description))
		}
		binary.LittleEndian.PutUint64(b[:], uint64(descOff))
		write(b[:], 8)
		binary.LittleEndian.PutUint32(b[:], uint32(len(r.Description)))
		write(b[:4], 4)
	}
	for _, d := range descs {
		write([]byte(d), len(d))
	}
	return
}

func NewFromPacked(b []byte) PackedIPList {
	return PackedIPList(b)
}

type PackedIPList []byte

var _ Ranger = PackedIPList{}

func (me PackedIPList) len() int {
	return int(binary.LittleEndian.Uint64(me[:8]))
}

func (me PackedIPList) NumRanges() int {
	return me.len()
}

func (me PackedIPList) getFirst(i int) net.IP {
	off := packedRangesOffset + packedRangeLen*i
	return net.IP(me[off : off+4])
}

func (me PackedIPList) getRange(i int) (ret Range) {
	rOff := packedRangesOffset + packedRangeLen*i
	last := me[rOff+4 : rOff+8]
	descOff := int(binary.LittleEndian.Uint64(me[rOff+8:]))
	descLen := int(binary.LittleEndian.Uint32(me[rOff+16:]))
	descOff += packedRangesOffset + packedRangeLen*me.len()
	ret = Range{
		me.getFirst(i),
		net.IP(last),
		string(me[descOff : descOff+descLen]),
	}
	return
}

func (me PackedIPList) Lookup(ip net.IP) (r Range, ok bool) {
	ip4 := ip.To4()
	if ip4 == nil {
		// If the IP list was built successfully, then it only contained IPv4
		// ranges. Therefore no IPv6 ranges are blocked.
		if ip.To16() == nil {
			r = Range{
				Description: "bad IP",
			}
			ok = true
		}
		return
	}
	return lookup(me.getFirst, me.getRange, me.len(), ip4)
}
<commit_msg>iplist: Add helper to mmap in a packed blocklist<commit_after>package iplist

import (
	"encoding/binary"
	"io"
	"net"
	"os"

	"github.com/edsrzf/mmap-go"
)

// The packed format is an 8 byte integer of the number of ranges. Then 20
// bytes per range, consisting of 4 byte packed IP being the lower bound IP of
// the range, then 4 bytes of the upper, inclusive bound, 8 bytes for the
// offset of the description from the end of the packed ranges, and 4 bytes
// for the length of the description. After these packed ranges, are the
// concatenated descriptions.

const (
	packedRangesOffset = 8
	packedRangeLen     = 20
)

func (me *IPList) WritePacked(w io.Writer) (err error) {
	descOffsets := make(map[string]int64, len(me.ranges))
	descs := make([]string, 0, len(me.ranges))
	var nextOffset int64
	// This is a little monadic, no?
	write := func(b []byte, expectedLen int) {
		if err != nil {
			return
		}
		var n int
		n, err = w.Write(b)
		if err != nil {
			return
		}
		if n != expectedLen {
			panic(n)
		}
	}
	var b [8]byte
	binary.LittleEndian.PutUint64(b[:], uint64(len(me.ranges)))
	write(b[:], 8)
	for _, r := range me.ranges {
		write(r.First.To4(), 4)
		write(r.Last.To4(), 4)
		descOff, ok := descOffsets[r.Description]
		if !ok {
			descOff = nextOffset
			descOffsets[r.Description] = descOff
			descs = append(descs, r.Description)
			nextOffset += int64(len(r.Description))
		}
		binary.LittleEndian.PutUint64(b[:], uint64(descOff))
		write(b[:], 8)
		binary.LittleEndian.PutUint32(b[:], uint32(len(r.Description)))
		write(b[:4], 4)
	}
	for _, d := range descs {
		write([]byte(d), len(d))
	}
	return
}

func NewFromPacked(b []byte) PackedIPList {
	return PackedIPList(b)
}

type PackedIPList []byte

var _ Ranger = PackedIPList{}

func (me PackedIPList) len() int {
	return int(binary.LittleEndian.Uint64(me[:8]))
}

func (me PackedIPList) NumRanges() int {
	return me.len()
}

func (me PackedIPList) getFirst(i int) net.IP {
	off := packedRangesOffset + packedRangeLen*i
	return net.IP(me[off : off+4])
}

func (me PackedIPList) getRange(i int) (ret Range) {
	rOff := packedRangesOffset + packedRangeLen*i
	last := me[rOff+4 : rOff+8]
	descOff := int(binary.LittleEndian.Uint64(me[rOff+8:]))
	descLen := int(binary.LittleEndian.Uint32(me[rOff+16:]))
	descOff += packedRangesOffset + packedRangeLen*me.len()
	ret = Range{
		me.getFirst(i),
		net.IP(last),
		string(me[descOff : descOff+descLen]),
	}
	return
}

func (me PackedIPList) Lookup(ip net.IP) (r Range, ok bool) {
	ip4 := ip.To4()
	if ip4 == nil {
		// If the IP list was built successfully, then it only contained IPv4
		// ranges. Therefore no IPv6 ranges are blocked.
		if ip.To16() == nil {
			r = Range{
				Description: "bad IP",
			}
			ok = true
		}
		return
	}
	return lookup(me.getFirst, me.getRange, me.len(), ip4)
}

func MMapPacked(filename string) (ret Ranger, err error) {
	f, err := os.Open(filename)
	if os.IsNotExist(err) {
		err = nil
		return
	}
	if err != nil {
		return
	}
	defer f.Close()
	mm, err := mmap.Map(f, mmap.RDONLY, 0)
	if err != nil {
		return
	}
	ret = NewFromPacked(mm)
	return
}
<|endoftext|>


<commit_before>// Datasource package contains database/source type related.  A few datasources
// are implemented here (test, csv).  This package also includes
// schema base services (datasource registry).
package datasource

import (
	"fmt"
	"strings"
	"sync"

	u "github.com/araddon/gou"

	"github.com/araddon/qlbridge/schema"
)

var (
	// the global data sources registry mutex
	registryMu sync.RWMutex
	// registry for data sources
	registry = newRegistry()

	// If disableRecover=true, we will not capture/suppress panics
	// Test only feature hopefully
	DisableRecover bool
)

// Register makes a datasource available by the provided @sourceName
// If Register is called twice with the same name or if source is nil, it panics.
//
//  Sources are specific schemas of type csv, elasticsearch, etc containing
//    multiple tables
func Register(sourceName string, source schema.Source) {
	registryMu.Lock()
	defer registryMu.Unlock()
	sourceName = strings.ToLower(sourceName)
	registerNeedsLock(sourceName, source)
}

func registerNeedsLock(sourceName string, source schema.Source) {
	if source == nil {
		panic("qlbridge/datasource: Register Source is nil")
	}

	if _, dupe := registry.sources[sourceName]; dupe {
		panic("qlbridge/datasource: Register called twice for source " + sourceName)
	}
	registry.sources[sourceName] = source
}

// Register makes a datasource available by the provided @sourceName
// If Register is called twice with the same name or if source is nil, it panics.
//
//  Sources are specific schemas of type csv, elasticsearch, etc containing
//    multiple tables
func RegisterSchemaSource(schema, sourceName string, source schema.Source) *schema.Schema {
	sourceName = strings.ToLower(sourceName)
	registryMu.Lock()
	defer registryMu.Unlock()
	registerNeedsLock(sourceName, source)
	s, _ := createSchema(sourceName)
	registry.schemas[sourceName] = s
	return s
}

// DataSourcesRegistry get access to the shared/global
// registry of all datasource implementations
func DataSourcesRegistry() *Registry {
	return registry
}

// Open a datasource, Global open connection function using
//  default schema registry
//
func OpenConn(sourceName, sourceConfig string) (schema.Conn, error) {
	sourcei, ok := registry.sources[sourceName]
	if !ok {
		return nil, fmt.Errorf("datasource: unknown source %q (forgotten import?)", sourceName)
	}
	source, err := sourcei.Open(sourceConfig)
	if err != nil {
		return nil, err
	}
	return source, nil
}

// Our internal map of different types of datasources that are registered
// for our runtime system to use
type Registry struct {
	// Map of source name, each source name is name of db in a specific source
	//   such as elasticsearch, mongo, csv etc
	sources map[string]schema.Source
	schemas map[string]*schema.Schema
	// We need to be able to flatten all tables across all sources into single keyspace
	//tableSources map[string]schema.DataSource
	tables []string
}

func newRegistry() *Registry {
	return &Registry{
		sources: make(map[string]schema.Source),
		schemas: make(map[string]*schema.Schema),
		//tableSources: make(map[string]schema.DataSource),
		tables: make([]string, 0),
	}
}

// Init pre-schema load call any sources that need pre-schema init
func (m *Registry) Init() {
	//registryMu.RLock()
	//defer registryMu.RUnlock()
	// TODO:  this is a race, we need a lock on sources
	for _, src := range m.sources {
		src.Init()
	}
}

// Get connection for given Database
//
//  @db      database name
//
func (m *Registry) Conn(db string) schema.Conn {

	//u.Debugf("Registry.Conn(db='%v') ", db)
	source := m.Get(strings.ToLower(db))
	if source != nil {
		//u.Debugf("found source: db=%s   %T", db, source)
		conn, err := source.Open(db)
		if err != nil {
			u.Errorf("could not open data source: %v  %v", db, err)
			return nil
		}
		//u.Infof("source: %T  %#v", conn, conn)
		return conn
	} else {
		u.Errorf("DataSource(%s) was not found", db)
	}
	return nil
}

// Get schema for given source
//
//  @schemaName =  virtual database name made up of multiple backend-sources
//
func (m *Registry) Schema(schemaName string) (*schema.Schema, bool) {

	registryMu.RLock()
	s, ok := m.schemas[schemaName]
	registryMu.RUnlock()
	if ok && s != nil {
		return s, ok
	}
	registryMu.Lock()
	defer registryMu.Unlock()
	s, ok = createSchema(schemaName)
	if ok {
		u.Debugf("s:%p datasource register schema %q", s, schemaName)
		m.schemas[schemaName] = s
	}
	return s, ok
}

// SchemaAdd Add a new Schema
func (m *Registry) SchemaAdd(s *schema.Schema) {
	registryMu.Lock()
	defer registryMu.Unlock()
	_, ok := m.schemas[s.Name]
	if ok {
		return
	}
	m.schemas[s.Name] = s
}

// Add a new SourceSchema to a schema which will be created if it doesn't exist
func (m *Registry) SourceSchemaAdd(schemaName string, ss *schema.SchemaSource) error {

	registryMu.RLock()
	s, ok := m.schemas[schemaName]
	registryMu.RUnlock()
	if !ok {
		u.Warnf("must have schema %#v", ss)
		return fmt.Errorf("Must have schema when adding source schema %v", ss.Name)
	}
	s.AddSourceSchema(ss)
	return loadSchema(ss)
}

// Schemas:  returns a list of schemas
func (m *Registry) Schemas() []string {

	registryMu.RLock()
	defer registryMu.RUnlock()
	schemas := make([]string, 0, len(m.schemas))
	for _, s := range m.schemas {
		schemas = append(m.schemas, s.Name)
	}
	return schemas
}

// Tables - Get all tables from this registry
func (m *Registry) Tables() []string {
	if len(m.tables) == 0 {
		tbls := make([]string, 0)
		for _, src := range m.sources {
			for _, tbl := range src.Tables() {
				tbls = append(tbls, tbl)
			}
		}
		m.tables = tbls
	}
	return m.tables
}

// given connection info, get datasource
//  @connInfo =    csv:///dev/stdin
//                 mockcsv
func (m *Registry) DataSource(connInfo string) schema.Source {
	// if  mysql.tablename allow that convention
	u.Debugf("get datasource: conn=%q ", connInfo)
	//parts := strings.SplitN(from, ".", 2)
	// TODO:  move this to a csv, or other source not in global registry
	sourceType := ""
	if len(connInfo) > 0 {
		switch {
		// case strings.HasPrefix(name, "file://"):
		// 	name = name[len("file://"):]
		case strings.HasPrefix(connInfo, "csv://"):
			sourceType = "csv"
			//m.db = connInfo[len("csv://"):]
		case strings.Contains(connInfo, "://"):
			strIdx := strings.Index(connInfo, "://")
			sourceType = connInfo[0:strIdx]
			//m.db = connInfo[strIdx+3:]
		default:
			sourceType = connInfo
		}
	}

	sourceType = strings.ToLower(sourceType)
	//u.Debugf("source: %v", sourceType)
	if source := m.Get(sourceType); source != nil {
		//u.Debugf("source: %T", source)
		return source
	} else {
		u.Errorf("DataSource(conn) was not found: '%v'", sourceType)
	}

	return nil
}

// Get a Data Source, similar to Source(@connInfo)
func (m *Registry) Get(sourceName string) schema.Source {
	return m.getDepth(0, sourceName)
}
func (m *Registry) getDepth(depth int, sourceName string) schema.Source {
	source, ok := m.sources[strings.ToLower(sourceName)]
	if ok {
		return source
	}
	if depth > 0 {
		return nil
	}
	parts := strings.SplitN(sourceName, "://", 2)
	if len(parts) == 2 {
		source = m.getDepth(1, parts[0])
		if source != nil {
			return source
		}
		u.Warnf("not able to find schema %q", sourceName)
	}
	return nil
}

func (m *Registry) String() string {
	sourceNames := make([]string, 0, len(m.sources))
	for source, _ := range m.sources {
		sourceNames = append(sourceNames, source)
	}
	return fmt.Sprintf("{Sources: [%s] }", strings.Join(sourceNames, ", "))
}

// Create a source schema from given named source
//  we will find Source for that name and introspect
func createSchema(sourceName string) (*schema.Schema, bool) {

	sourceName = strings.ToLower(sourceName)

	ss := schema.NewSchemaSource(sourceName, sourceName)
	//u.Debugf("ss:%p createSchema %v", ss, sourceName)

	ds := registry.Get(sourceName)
	if ds == nil {
		u.Warnf("not able to find schema %q", sourceName)
		return nil, false
	}

	ss.DS = ds
	s := schema.NewSchema(sourceName)
	s.AddSourceSchema(ss)
	loadSchema(ss)

	return s, true
}

func loadSchema(ss *schema.SchemaSource) error {

	if ss.DS == nil {
		u.Warnf("missing DataSource for %s", ss.Name)
		panic(fmt.Sprintf("Missing datasource for %q", ss.Name))
	}

	if dsConfig, getsConfig := ss.DS.(schema.SourceSetup); getsConfig {
		if err := dsConfig.Setup(ss); err != nil {
			u.Errorf("Error setuping up %v  %v", ss.Name, err)
			return err
		}
	}

	s := ss.Schema()
	infoSchema := s.InfoSchema
	var infoSchemaSource *schema.SchemaSource
	var err error

	if infoSchema == nil {

		infoSchema = schema.NewSchema("schema")
		infoSchemaSource = schema.NewSchemaSource("schema", "schema")

		schemaDb := NewSchemaDb(s)
		infoSchemaSource.DS = schemaDb
		schemaDb.is = infoSchema
		//u.Debugf("schema:%p ss:%p loadSystemSchema: NEW infoschema:%p  s:%s ss:%s", s, ss, infoSchema, s.Name, ss.Name)

		infoSchemaSource.AddTableName("tables")
		infoSchema.InfoSchema = infoSchema
		infoSchema.AddSourceSchema(infoSchemaSource)
	} else {
		infoSchemaSource, err = infoSchema.Source("schema")
	}

	if err != nil {
		u.Errorf("could not find schema")
		return err
	}

	// For each table in source schema
	for _, tableName := range ss.Tables() {
		//u.Debugf("adding table: %q to infoSchema %p", tableName, infoSchema)
		_, err := ss.Table(tableName)
		if err != nil {
			//u.Warnf("Missing table?? %q", tableName)
			continue
		}
		infoSchemaSource.AddTableName(tableName)
	}

	s.InfoSchema = infoSchema

	s.RefreshSchema()

	//u.Debugf("s:%p ss:%p infoschema:%p  name:%s", s, ss, infoSchema, s.Name)

	return nil
}
<commit_msg>Expose list of schemas<commit_after>// Datasource package contains database/source type related.  A few datasources
// are implemented here (test, csv).  This package also includes
// schema base services (datasource registry).
package datasource

import (
	"fmt"
	"strings"
	"sync"

	u "github.com/araddon/gou"

	"github.com/araddon/qlbridge/schema"
)

var (
	// the global data sources registry mutex
	registryMu sync.RWMutex
	// registry for data sources
	registry = newRegistry()

	// If disableRecover=true, we will not capture/suppress panics
	// Test only feature hopefully
	DisableRecover bool
)

// Register makes a datasource available by the provided @sourceName
// If Register is called twice with the same name or if source is nil, it panics.
//
//  Sources are specific schemas of type csv, elasticsearch, etc containing
//    multiple tables
func Register(sourceName string, source schema.Source) {
	registryMu.Lock()
	defer registryMu.Unlock()
	sourceName = strings.ToLower(sourceName)
	registerNeedsLock(sourceName, source)
}

func registerNeedsLock(sourceName string, source schema.Source) {
	if source == nil {
		panic("qlbridge/datasource: Register Source is nil")
	}

	if _, dupe := registry.sources[sourceName]; dupe {
		panic("qlbridge/datasource: Register called twice for source " + sourceName)
	}
	registry.sources[sourceName] = source
}

// Register makes a datasource available by the provided @sourceName
// If Register is called twice with the same name or if source is nil, it panics.
//
//  Sources are specific schemas of type csv, elasticsearch, etc containing
//    multiple tables
func RegisterSchemaSource(schema, sourceName string, source schema.Source) *schema.Schema {
	sourceName = strings.ToLower(sourceName)
	registryMu.Lock()
	defer registryMu.Unlock()
	registerNeedsLock(sourceName, source)
	s, _ := createSchema(sourceName)
	registry.schemas[sourceName] = s
	return s
}

// DataSourcesRegistry get access to the shared/global
// registry of all datasource implementations
func DataSourcesRegistry() *Registry {
	return registry
}

// Open a datasource, Global open connection function using
//  default schema registry
//
func OpenConn(sourceName, sourceConfig string) (schema.Conn, error) {
	sourcei, ok := registry.sources[sourceName]
	if !ok {
		return nil, fmt.Errorf("datasource: unknown source %q (forgotten import?)", sourceName)
	}
	source, err := sourcei.Open(sourceConfig)
	if err != nil {
		return nil, err
	}
	return source, nil
}

// Our internal map of different types of datasources that are registered
// for our runtime system to use
type Registry struct {
	// Map of source name, each source name is name of db in a specific source
	//   such as elasticsearch, mongo, csv etc
	sources map[string]schema.Source
	schemas map[string]*schema.Schema
	// We need to be able to flatten all tables across all sources into single keyspace
	//tableSources map[string]schema.DataSource
	tables []string
}

func newRegistry() *Registry {
	return &Registry{
		sources: make(map[string]schema.Source),
		schemas: make(map[string]*schema.Schema),
		//tableSources: make(map[string]schema.DataSource),
		tables: make([]string, 0),
	}
}

// Init pre-schema load call any sources that need pre-schema init
func (m *Registry) Init() {
	//registryMu.RLock()
	//defer registryMu.RUnlock()
	// TODO:  this is a race, we need a lock on sources
	for _, src := range m.sources {
		src.Init()
	}
}

// Get connection for given Database
//
//  @db      database name
//
func (m *Registry) Conn(db string) schema.Conn {

	//u.Debugf("Registry.Conn(db='%v') ", db)
	source := m.Get(strings.ToLower(db))
	if source != nil {
		//u.Debugf("found source: db=%s   %T", db, source)
		conn, err := source.Open(db)
		if err != nil {
			u.Errorf("could not open data source: %v  %v", db, err)
			return nil
		}
		//u.Infof("source: %T  %#v", conn, conn)
		return conn
	} else {
		u.Errorf("DataSource(%s) was not found", db)
	}
	return nil
}

// Get schema for given source
//
//  @schemaName =  virtual database name made up of multiple backend-sources
//
func (m *Registry) Schema(schemaName string) (*schema.Schema, bool) {

	registryMu.RLock()
	s, ok := m.schemas[schemaName]
	registryMu.RUnlock()
	if ok && s != nil {
		return s, ok
	}
	registryMu.Lock()
	defer registryMu.Unlock()
	s, ok = createSchema(schemaName)
	if ok {
		u.Debugf("s:%p datasource register schema %q", s, schemaName)
		m.schemas[schemaName] = s
	}
	return s, ok
}

// SchemaAdd Add a new Schema
func (m *Registry) SchemaAdd(s *schema.Schema) {
	registryMu.Lock()
	defer registryMu.Unlock()
	_, ok := m.schemas[s.Name]
	if ok {
		return
	}
	m.schemas[s.Name] = s
}

// Add a new SourceSchema to a schema which will be created if it doesn't exist
func (m *Registry) SourceSchemaAdd(schemaName string, ss *schema.SchemaSource) error {

	registryMu.RLock()
	s, ok := m.schemas[schemaName]
	registryMu.RUnlock()
	if !ok {
		u.Warnf("must have schema %#v", ss)
		return fmt.Errorf("Must have schema when adding source schema %v", ss.Name)
	}
	s.AddSourceSchema(ss)
	return loadSchema(ss)
}

// Schemas:  returns a list of schemas
func (m *Registry) Schemas() []string {

	registryMu.RLock()
	defer registryMu.RUnlock()
	schemas := make([]string, 0, len(m.schemas))
	for _, s := range m.schemas {
		schemas = append(schemas, s.Name)
	}
	return schemas
}

// Tables - Get all tables from this registry
func (m *Registry) Tables() []string {
	if len(m.tables) == 0 {
		tbls := make([]string, 0)
		for _, src := range m.sources {
			for _, tbl := range src.Tables() {
				tbls = append(tbls, tbl)
			}
		}
		m.tables = tbls
	}
	return m.tables
}

// given connection info, get datasource
//  @connInfo =    csv:///dev/stdin
//                 mockcsv
func (m *Registry) DataSource(connInfo string) schema.Source {
	// if  mysql.tablename allow that convention
	u.Debugf("get datasource: conn=%q ", connInfo)
	//parts := strings.SplitN(from, ".", 2)
	// TODO:  move this to a csv, or other source not in global registry
	sourceType := ""
	if len(connInfo) > 0 {
		switch {
		// case strings.HasPrefix(name, "file://"):
		// 	name = name[len("file://"):]
		case strings.HasPrefix(connInfo, "csv://"):
			sourceType = "csv"
			//m.db = connInfo[len("csv://"):]
		case strings.Contains(connInfo, "://"):
			strIdx := strings.Index(connInfo, "://")
			sourceType = connInfo[0:strIdx]
			//m.db = connInfo[strIdx+3:]
		default:
			sourceType = connInfo
		}
	}

	sourceType = strings.ToLower(sourceType)
	//u.Debugf("source: %v", sourceType)
	if source := m.Get(sourceType); source != nil {
		//u.Debugf("source: %T", source)
		return source
	} else {
		u.Errorf("DataSource(conn) was not found: '%v'", sourceType)
	}

	return nil
}

// Get a Data Source, similar to Source(@connInfo)
func (m *Registry) Get(sourceName string) schema.Source {
	return m.getDepth(0, sourceName)
}
func (m *Registry) getDepth(depth int, sourceName string) schema.Source {
	source, ok := m.sources[strings.ToLower(sourceName)]
	if ok {
		return source
	}
	if depth > 0 {
		return nil
	}
	parts := strings.SplitN(sourceName, "://", 2)
	if len(parts) == 2 {
		source = m.getDepth(1, parts[0])
		if source != nil {
			return source
		}
		u.Warnf("not able to find schema %q", sourceName)
	}
	return nil
}

func (m *Registry) String() string {
	sourceNames := make([]string, 0, len(m.sources))
	for source, _ := range m.sources {
		sourceNames = append(sourceNames, source)
	}
	return fmt.Sprintf("{Sources: [%s] }", strings.Join(sourceNames, ", "))
}

// Create a source schema from given named source
//  we will find Source for that name and introspect
func createSchema(sourceName string) (*schema.Schema, bool) {

	sourceName = strings.ToLower(sourceName)

	ss := schema.NewSchemaSource(sourceName, sourceName)
	//u.Debugf("ss:%p createSchema %v", ss, sourceName)

	ds := registry.Get(sourceName)
	if ds == nil {
		u.Warnf("not able to find schema %q", sourceName)
		return nil, false
	}

	ss.DS = ds
	s := schema.NewSchema(sourceName)
	s.AddSourceSchema(ss)
	loadSchema(ss)

	return s, true
}

func loadSchema(ss *schema.SchemaSource) error {

	if ss.DS == nil {
		u.Warnf("missing DataSource for %s", ss.Name)
		panic(fmt.Sprintf("Missing datasource for %q", ss.Name))
	}

	if dsConfig, getsConfig := ss.DS.(schema.SourceSetup); getsConfig {
		if err := dsConfig.Setup(ss); err != nil {
			u.Errorf("Error setuping up %v  %v", ss.Name, err)
			return err
		}
	}

	s := ss.Schema()
	infoSchema := s.InfoSchema
	var infoSchemaSource *schema.SchemaSource
	var err error

	if infoSchema == nil {

		infoSchema = schema.NewSchema("schema")
		infoSchemaSource = schema.NewSchemaSource("schema", "schema")

		schemaDb := NewSchemaDb(s)
		infoSchemaSource.DS = schemaDb
		schemaDb.is = infoSchema
		//u.Debugf("schema:%p ss:%p loadSystemSchema: NEW infoschema:%p  s:%s ss:%s", s, ss, infoSchema, s.Name, ss.Name)

		infoSchemaSource.AddTableName("tables")
		infoSchema.InfoSchema = infoSchema
		infoSchema.AddSourceSchema(infoSchemaSource)
	} else {
		infoSchemaSource, err = infoSchema.Source("schema")
	}

	if err != nil {
		u.Errorf("could not find schema")
		return err
	}

	// For each table in source schema
	for _, tableName := range ss.Tables() {
		//u.Debugf("adding table: %q to infoSchema %p", tableName, infoSchema)
		_, err := ss.Table(tableName)
		if err != nil {
			//u.Warnf("Missing table?? %q", tableName)
			continue
		}
		infoSchemaSource.AddTableName(tableName)
	}

	s.InfoSchema = infoSchema

	s.RefreshSchema()

	//u.Debugf("s:%p ss:%p infoschema:%p  name:%s", s, ss, infoSchema, s.Name)

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"os"
	"sync"

	"github.com/anacrolix/log"

	_ "github.com/anacrolix/envpprof"
	"github.com/anacrolix/tagflag"

	"github.com/anacrolix/dht/v2"
)

func main() {
	code := mainErr()
	if code != 0 {
		os.Exit(code)
	}
}

func mainErr() int {
	var flags = struct {
		Port  int
		Debug bool
		tagflag.StartPos
		Infohash [][20]byte
	}{}
	tagflag.Parse(&flags)
	s, err := dht.NewServer(func() *dht.ServerConfig {
		sc := dht.NewDefaultServerConfig()
		if flags.Debug {
			sc.Logger = log.Default
		}
		return sc
	}())
	if err != nil {
		log.Printf("error creating server: %s", err)
		return 1
	}
	defer s.Close()
	wg := sync.WaitGroup{}
	addrs := make(map[[20]byte]map[string]struct{}, len(flags.Infohash))
	for _, ih := range flags.Infohash {
		a, err := s.Announce(ih, flags.Port, false)
		if err != nil {
			log.Printf("error announcing %s: %s", ih, err)
			continue
		}
		wg.Add(1)
		addrs[ih] = make(map[string]struct{})
		go func(ih [20]byte) {
			defer wg.Done()
			for ps := range a.Peers {
				for _, p := range ps.Peers {
					s := p.String()
					if _, ok := addrs[ih][s]; !ok {
						log.Printf("got peer %s for %x from %s", p, ih, ps.NodeInfo)
						addrs[ih][s] = struct{}{}
					}
				}
			}
			log.Printf("%v contacted %v nodes", a, a.NumContacted())
		}(ih)
	}
	wg.Wait()
	for _, ih := range flags.Infohash {
		ips := make(map[string]struct{}, len(addrs[ih]))
		for s := range addrs[ih] {
			ip, _, err := net.SplitHostPort(s)
			if err != nil {
				log.Printf("error parsing addr: %s", err)
			}
			ips[ip] = struct{}{}
		}
		log.Printf("%x: %d addrs %d distinct ips", ih, len(addrs[ih]), len(ips))
	}
	return 0
}
<commit_msg>cmd/dht-announce: Close gracefully on interrupt<commit_after>package main

import (
	"net"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/anacrolix/log"

	_ "github.com/anacrolix/envpprof"
	"github.com/anacrolix/tagflag"

	"github.com/anacrolix/dht/v2"
)

func main() {
	code := mainErr()
	if code != 0 {
		os.Exit(code)
	}
}

func mainErr() int {
	var flags = struct {
		Port  int
		Debug bool
		tagflag.StartPos
		Infohash [][20]byte
	}{}
	tagflag.Parse(&flags)
	s, err := dht.NewServer(func() *dht.ServerConfig {
		sc := dht.NewDefaultServerConfig()
		if flags.Debug {
			sc.Logger = log.Default
		}
		return sc
	}())
	if err != nil {
		log.Printf("error creating server: %s", err)
		return 1
	}
	defer s.Close()
	var wg sync.WaitGroup
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	stop := make(chan struct{})
	go func() {
		<-sigChan
		close(stop)
	}()
	addrs := make(map[[20]byte]map[string]struct{}, len(flags.Infohash))
	for _, ih := range flags.Infohash {
		a, err := s.Announce(ih, flags.Port, false)
		if err != nil {
			log.Printf("error announcing %s: %s", ih, err)
			continue
		}
		wg.Add(1)
		addrs[ih] = make(map[string]struct{})
		go func(ih [20]byte) {
			defer wg.Done()
			for ps := range a.Peers {
				for _, p := range ps.Peers {
					s := p.String()
					if _, ok := addrs[ih][s]; !ok {
						log.Printf("got peer %s for %x from %s", p, ih, ps.NodeInfo)
						addrs[ih][s] = struct{}{}
					}
				}
			}
			time.Sleep(time.Second)
			log.Printf("%v contacted %v nodes", a, a.NumContacted())
		}(ih)
		go func() {
			<-stop
			a.Close()
		}()
	}
	wg.Wait()
	for _, ih := range flags.Infohash {
		ips := make(map[string]struct{}, len(addrs[ih]))
		for s := range addrs[ih] {
			ip, _, err := net.SplitHostPort(s)
			if err != nil {
				log.Printf("error parsing addr: %s", err)
			}
			ips[ip] = struct{}{}
		}
		log.Printf("%x: %d addrs %d distinct ips", ih, len(addrs[ih]), len(ips))
	}
	return 0
}
<|endoftext|>


<commit_before>package driver

import (
	"io"
	"log"
	"net"
	"net/rpc"

	"github.com/hashicorp/go-plugin"
	"github.com/hashicorp/nomad/client/driver/executor"
)

var HandshakeConfig = plugin.HandshakeConfig{
	ProtocolVersion:  1,
	MagicCookieKey:   "NOMAD_PLUGIN_MAGIC_COOKIE",
	MagicCookieValue: "e4327c2e01eabfd75a8a67adb114fb34a757d57eee7728d857a8cec6e91a7255",
}

func GetPluginMap(w io.Writer) map[string]plugin.Plugin {
	p := new(ExecutorPlugin)
	p.logger = log.New(w, "executor-plugin-server:", log.LstdFlags)
	return map[string]plugin.Plugin{"executor": p}
}

// ExecutorReattachConfig is the config that we seralize and de-serialize and
// store in disk
type ExecutorReattachConfig struct {
	Pid      int
	AddrNet  string
	AddrName string
}

// PluginConfig returns a config from an ExecutorReattachConfig
func (c *ExecutorReattachConfig) PluginConfig() *plugin.ReattachConfig {
	var addr net.Addr
	switch c.AddrNet {
	case "unix", "unixgram", "unixpacket":
		addr, _ = net.ResolveUnixAddr(c.AddrNet, c.AddrName)
	case "tcp", "tcp4", "tcp6":
		addr, _ = net.ResolveTCPAddr(c.AddrNet, c.AddrName)
	}
	return &plugin.ReattachConfig{Pid: c.Pid, Addr: addr}
}

func NewExecutorReattachConfig(c *plugin.ReattachConfig) *ExecutorReattachConfig {
	return &ExecutorReattachConfig{Pid: c.Pid, AddrNet: c.Addr.Network(), AddrName: c.Addr.String()}
}

type ExecutorRPC struct {
	client *rpc.Client
}

// LaunchCmdArgs wraps a user command and the args for the purposes of RPC
type LaunchCmdArgs struct {
	Cmd *executor.ExecCommand
	Ctx *executor.ExecutorContext
}

func (e *ExecutorRPC) LaunchCmd(cmd *executor.ExecCommand, ctx *executor.ExecutorContext) (*executor.ProcessState, error) {
	var ps *executor.ProcessState
	err := e.client.Call("Plugin.LaunchCmd", LaunchCmdArgs{Cmd: cmd, Ctx: ctx}, &ps)
	return ps, err
}

func (e *ExecutorRPC) Wait() (*executor.ProcessState, error) {
	var ps executor.ProcessState
	err := e.client.Call("Plugin.Wait", new(interface{}), &ps)
	return &ps, err
}

func (e *ExecutorRPC) ShutDown() error {
	return e.client.Call("Plugin.ShutDown", new(interface{}), new(interface{}))
}

func (e *ExecutorRPC) Exit() error {
	return e.client.Call("Plugin.Exit", new(interface{}), new(interface{}))
}

type ExecutorRPCServer struct {
	Impl executor.Executor
}

func (e *ExecutorRPCServer) LaunchCmd(args LaunchCmdArgs, ps *executor.ProcessState) error {
	state, err := e.Impl.LaunchCmd(args.Cmd, args.Ctx)
	if state != nil {
		*ps = *state
	}
	return err
}

func (e *ExecutorRPCServer) Wait(args interface{}, ps *executor.ProcessState) error {
	state, err := e.Impl.Wait()
	if state != nil {
		*ps = *state
	}
	return err
}

func (e *ExecutorRPCServer) ShutDown(args interface{}, resp *interface{}) error {
	return e.Impl.ShutDown()
}

func (e *ExecutorRPCServer) Exit(args interface{}, resp *interface{}) error {
	return e.Impl.Exit()
}

type ExecutorPlugin struct {
	logger *log.Logger
	Impl   *ExecutorRPCServer
}

func (p *ExecutorPlugin) Server(*plugin.MuxBroker) (interface{}, error) {
	if p.Impl == nil {
		p.Impl = &ExecutorRPCServer{Impl: executor.NewExecutor(p.logger)}
	}
	return p.Impl, nil
}

func (p *ExecutorPlugin) Client(b *plugin.MuxBroker, c *rpc.Client) (interface{}, error) {
	return &ExecutorRPC{client: c}, nil
}
<commit_msg>removing the prefix of the logger<commit_after>package driver

import (
	"io"
	"log"
	"net"
	"net/rpc"

	"github.com/hashicorp/go-plugin"
	"github.com/hashicorp/nomad/client/driver/executor"
)

var HandshakeConfig = plugin.HandshakeConfig{
	ProtocolVersion:  1,
	MagicCookieKey:   "NOMAD_PLUGIN_MAGIC_COOKIE",
	MagicCookieValue: "e4327c2e01eabfd75a8a67adb114fb34a757d57eee7728d857a8cec6e91a7255",
}

func GetPluginMap(w io.Writer) map[string]plugin.Plugin {
	p := new(ExecutorPlugin)
	p.logger = log.New(w, "", log.LstdFlags)
	return map[string]plugin.Plugin{"executor": p}
}

// ExecutorReattachConfig is the config that we seralize and de-serialize and
// store in disk
type ExecutorReattachConfig struct {
	Pid      int
	AddrNet  string
	AddrName string
}

// PluginConfig returns a config from an ExecutorReattachConfig
func (c *ExecutorReattachConfig) PluginConfig() *plugin.ReattachConfig {
	var addr net.Addr
	switch c.AddrNet {
	case "unix", "unixgram", "unixpacket":
		addr, _ = net.ResolveUnixAddr(c.AddrNet, c.AddrName)
	case "tcp", "tcp4", "tcp6":
		addr, _ = net.ResolveTCPAddr(c.AddrNet, c.AddrName)
	}
	return &plugin.ReattachConfig{Pid: c.Pid, Addr: addr}
}

func NewExecutorReattachConfig(c *plugin.ReattachConfig) *ExecutorReattachConfig {
	return &ExecutorReattachConfig{Pid: c.Pid, AddrNet: c.Addr.Network(), AddrName: c.Addr.String()}
}

type ExecutorRPC struct {
	client *rpc.Client
}

// LaunchCmdArgs wraps a user command and the args for the purposes of RPC
type LaunchCmdArgs struct {
	Cmd *executor.ExecCommand
	Ctx *executor.ExecutorContext
}

func (e *ExecutorRPC) LaunchCmd(cmd *executor.ExecCommand, ctx *executor.ExecutorContext) (*executor.ProcessState, error) {
	var ps *executor.ProcessState
	err := e.client.Call("Plugin.LaunchCmd", LaunchCmdArgs{Cmd: cmd, Ctx: ctx}, &ps)
	return ps, err
}

func (e *ExecutorRPC) Wait() (*executor.ProcessState, error) {
	var ps executor.ProcessState
	err := e.client.Call("Plugin.Wait", new(interface{}), &ps)
	return &ps, err
}

func (e *ExecutorRPC) ShutDown() error {
	return e.client.Call("Plugin.ShutDown", new(interface{}), new(interface{}))
}

func (e *ExecutorRPC) Exit() error {
	return e.client.Call("Plugin.Exit", new(interface{}), new(interface{}))
}

type ExecutorRPCServer struct {
	Impl executor.Executor
}

func (e *ExecutorRPCServer) LaunchCmd(args LaunchCmdArgs, ps *executor.ProcessState) error {
	state, err := e.Impl.LaunchCmd(args.Cmd, args.Ctx)
	if state != nil {
		*ps = *state
	}
	return err
}

func (e *ExecutorRPCServer) Wait(args interface{}, ps *executor.ProcessState) error {
	state, err := e.Impl.Wait()
	if state != nil {
		*ps = *state
	}
	return err
}

func (e *ExecutorRPCServer) ShutDown(args interface{}, resp *interface{}) error {
	return e.Impl.ShutDown()
}

func (e *ExecutorRPCServer) Exit(args interface{}, resp *interface{}) error {
	return e.Impl.Exit()
}

type ExecutorPlugin struct {
	logger *log.Logger
	Impl   *ExecutorRPCServer
}

func (p *ExecutorPlugin) Server(*plugin.MuxBroker) (interface{}, error) {
	if p.Impl == nil {
		p.Impl = &ExecutorRPCServer{Impl: executor.NewExecutor(p.logger)}
	}
	return p.Impl, nil
}

func (p *ExecutorPlugin) Client(b *plugin.MuxBroker, c *rpc.Client) (interface{}, error) {
	return &ExecutorRPC{client: c}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"runtime"
	"strings"

	"github.com/pkg/errors"
	"github.com/urfave/cli"
)

const (
	caddyfile = `
0.0.0.0:3000 {
    proxy / localhost:3010
    proxy /api localhost:3020 {
        without /api
    }
}
`
)

func ActionDevSetup(c *cli.Context) error {
	if err := os.MkdirAll("./dev", 0755); err != nil {
		return errors.Wrap(err, "failed to create ./dev/ for development")
	}
	log.Println("Created ./dev/")

	if err := setupCaddy(); err != nil {
		return err
	}

	return nil
}

func setupCaddy() error {
	if err := downloadCaddy(); err != nil {
		return errors.Wrap(err, "failed to download caddy")
	}
	log.Println("Downloaded ./dev/caddy.zip")

	if err := extractCaddy(); err != nil {
		return errors.Wrap(err, "failed to extract caddy")
	}
	log.Println("Extracted ./dev/caddy.zip to ./dev/caddy")

	// Create Caddyfile with contents if it not exist
	if err := createCaddyfile(); err != nil {
		return errors.Wrap(err, "failed to create ./dev/Caddyfile")
	}
	log.Println("Created ./dev/Caddyfile")

	return nil
}

func downloadCaddy() error {
	caddyURL := fmt.Sprintf("https://caddyserver.com/download/%s/%s", runtime.GOOS, runtime.GOARCH)

	// Download & extract Caddy to ./dev/caddy if it not exist

	exist, err := exists("./dev/caddy")
	if err != nil {
		return err
	}
	if exist { // if it exist, don't do more
		return nil
	}

	exist, err = exists("./dev/caddy.zip")
	if err != nil {
		return err
	}
	if exist { // if it exist, don't do more
		return nil
	}

	out, err := os.Create("./dev/caddy.zip")
	if err != nil {
		return err
	}
	defer out.Close()

	resp, err := http.Get(caddyURL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	_, err = io.Copy(out, resp.Body)
	if err != nil {
		return err
	}

	return nil
}

func extractCaddy() error {
	r, err := zip.OpenReader("./dev/caddy.zip")
	if err != nil {
		return err
	}

	var zippedCaddy *zip.File
	for _, file := range r.File {
		if file.Name == "caddy" {
			zippedCaddy = file
		}
	}

	fileReader, err := zippedCaddy.Open()
	if err != nil {
		return err
	}
	defer fileReader.Close()

	targetFile, err := os.OpenFile("./dev/caddy", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, zippedCaddy.Mode())
	if err != nil {
		return err
	}
	defer targetFile.Close()

	if _, err := io.Copy(targetFile, fileReader); err != nil {
		return err
	}

	return nil
}

func createCaddyfile() error {
	exist, err := exists("./dev/Caddyfile")
	if err != nil {
		return err
	}
	if !exist {
		if err := ioutil.WriteFile("./dev/Caddyfile", []byte(strings.TrimSpace(caddyfile)), 0644); err != nil {
			return err
		}
	}

	return nil
}

func exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return true, err
}
<commit_msg>If GOOS==darwin unzip caddy otherwise extract .tar.gz<commit_after>package main

import (
	"archive/tar"
	"archive/zip"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"runtime"
	"strings"

	"github.com/pkg/errors"
	"github.com/urfave/cli"
)

const (
	caddyfile = `
0.0.0.0:3000 {
    proxy / localhost:3010
    proxy /api localhost:3020 {
        without /api
    }
}
`
)

func ActionDevSetup(c *cli.Context) error {
	if err := os.MkdirAll("./dev", 0755); err != nil {
		return errors.Wrap(err, "failed to create ./dev/ for development")
	}
	log.Println("Created ./dev/")

	if err := setupCaddy(); err != nil {
		return err
	}

	return nil
}

func setupCaddy() error {
	url := fmt.Sprintf("https://caddyserver.com/download/%s/%s", runtime.GOOS, runtime.GOARCH)
	archive := ""

	switch runtime.GOOS {
	case "darwin":
		archive = "./dev/caddy.zip"
	default:
		archive = "./dev/caddy.tar.gz"
	}

	if err := downloadCaddy(url, archive); err != nil {
		return errors.Wrap(err, "failed to download caddy")
	}
	log.Println("Downloaded ./dev/caddy.zip")

	if err := extractCaddy(archive); err != nil {
		return errors.Wrap(err, "failed to extract caddy")
	}
	log.Println("Extracted ./dev/caddy.zip to ./dev/caddy")

	// Create Caddyfile with contents if it not exist
	if err := createCaddyfile(); err != nil {
		return errors.Wrap(err, "failed to create ./dev/Caddyfile")
	}
	log.Println("Created ./dev/Caddyfile")

	return nil
}

func downloadCaddy(url, archive string) error {

	// Download & extract Caddy to ./dev/caddy if it not exist

	exist, err := exists("./dev/caddy")
	if err != nil {
		return err
	}
	if exist { // if it exist, don't do more
		return nil
	}

	exist, err = exists(archive)
	if err != nil {
		return err
	}
	if exist { // if it exist, don't do more
		return nil
	}

	out, err := os.Create(archive)
	if err != nil {
		return err
	}
	defer out.Close()

	log.Println("Downloading", url)
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	_, err = io.Copy(out, resp.Body)
	if err != nil {
		return err
	}

	return nil
}

func extractCaddy(archive string) error {
	switch runtime.GOOS {
	case "darwin":
		r, err := zip.OpenReader(archive)
		if err != nil {
			return err
		}

		var zippedCaddy *zip.File
		for _, file := range r.File {
			if file.Name == "caddy" {
				zippedCaddy = file
			}
		}

		fileReader, err := zippedCaddy.Open()
		if err != nil {
			return err
		}
		defer fileReader.Close()

		targetFile, err := os.OpenFile("./dev/caddy", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, zippedCaddy.Mode())
		if err != nil {
			return err
		}
		defer targetFile.Close()

		if _, err := io.Copy(targetFile, fileReader); err != nil {
			return err
		}
	default:
		archiveFile, err := os.Open(archive)
		if err != nil {
			return err
		}
		defer archiveFile.Close()

		gzipREader, err := gzip.NewReader(archiveFile)
		if err != nil {
			return err
		}

		tarReader := tar.NewReader(gzipREader)
		for {
			header, err := tarReader.Next()

			if err == io.EOF {
				break
			}

			if err != nil {
				return err
			}

			if header.Name == "caddy" && header.Typeflag == tar.TypeReg {
				f, err := os.Create("./dev/caddy")
				if err != nil {
					return err
				}
				defer f.Close()

				if err := f.Chmod(0744); err != nil {
					return err
				}

				if _, err := io.Copy(f, tarReader); err != nil {
					return err
				}
				return nil
			}
		}
	}
	return nil
}

func createCaddyfile() error {
	exist, err := exists("./dev/Caddyfile")
	if err != nil {
		return err
	}
	if !exist {
		if err := ioutil.WriteFile("./dev/Caddyfile", []byte(strings.TrimSpace(caddyfile)), 0644); err != nil {
			return err
		}
	}

	return nil
}

func exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return true, err
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gocli"
)

type Whoami struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Whoami) Run(args []string) (exitCode int) {
	var tags = map[string]string{
		// pubsub
		"10.209.36.14": "psub 1",
		"10.209.36.33": "psub 2",
		"10.213.1.210": "psub 3",
		"10.213.9.245": "psub 4",

		// kafka brokers with small disks
		"10.209.37.39":  "kk 1",
		"10.209.33.20":  "kk 2",
		"10.209.37.69":  "kk 3",
		"10.209.33.40":  "kk 4",
		"10.209.11.166": "kk 5",
		"10.209.11.195": "kk 6",
		"10.209.10.161": "kk 7",
		"10.209.10.141": "kk 8",

		// kafka brokers with big disks
		"10.209.18.15":   "k 1",
		"10.209.18.16":   "k 2",
		"10.209.18.65":   "k 3",
		"10.209.18.66":   "k 4",
		"10.209.19.143":  "k 5",
		"10.209.19.144":  "k 6",
		"10.209.22.142":  "k 7",
		"10.209.19.35":   "k 8",
		"10.209.19.36":   "k 9",
		"10.209.240.191": "k 11",
		"10.209.240.192": "k 12",
		"10.209.240.193": "k 13",
		"10.209.240.194": "k 14",

		// zk
		"10.209.33.69": "czk 1",
		"10.209.37.19": "czk 2",
		"10.209.37.68": "czk 3",
	}

	ip, _ := ctx.LocalIP()
	if tag, present := tags[ip.String()]; present {
		this.Ui.Output(tag)
	} else {
		this.Ui.Warn("unknown")
	}

	return
}

func (*Whoami) Synopsis() string {
	return "Display effective expect script tag of current host"
}

func (this *Whoami) Help() string {
	help := fmt.Sprintf(`
Usage: %s whoami [options]

    Display effective expect script tag of current host

`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>add a new whoami host<commit_after>package command

import (
	"fmt"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gocli"
)

type Whoami struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Whoami) Run(args []string) (exitCode int) {
	var tags = map[string]string{
		// misc
		"10.213.57.149": "dev",

		// pubsub
		"10.209.36.14": "psub 1",
		"10.209.36.33": "psub 2",
		"10.213.1.210": "psub 3",
		"10.213.9.245": "psub 4",

		// kafka brokers with small disks
		"10.209.37.39":  "kk 1",
		"10.209.33.20":  "kk 2",
		"10.209.37.69":  "kk 3",
		"10.209.33.40":  "kk 4",
		"10.209.11.166": "kk 5",
		"10.209.11.195": "kk 6",
		"10.209.10.161": "kk 7",
		"10.209.10.141": "kk 8",

		// kafka brokers with big disks
		"10.209.18.15":   "k 1",
		"10.209.18.16":   "k 2",
		"10.209.18.65":   "k 3",
		"10.209.18.66":   "k 4",
		"10.209.19.143":  "k 5",
		"10.209.19.144":  "k 6",
		"10.209.22.142":  "k 7",
		"10.209.19.35":   "k 8",
		"10.209.19.36":   "k 9",
		"10.209.240.191": "k 11",
		"10.209.240.192": "k 12",
		"10.209.240.193": "k 13",
		"10.209.240.194": "k 14",

		// zk
		"10.209.33.69": "czk 1",
		"10.209.37.19": "czk 2",
		"10.209.37.68": "czk 3",
	}

	ip, _ := ctx.LocalIP()
	if tag, present := tags[ip.String()]; present {
		this.Ui.Output(tag)
	} else {
		this.Ui.Warn("unknown")
	}

	return
}

func (*Whoami) Synopsis() string {
	return "Display effective expect script tag of current host"
}

func (this *Whoami) Help() string {
	help := fmt.Sprintf(`
Usage: %s whoami [options]

    Display effective expect script tag of current host

`, this.Cmd)
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package mc

import (
	"os"
	"path/filepath"
	"strings"

	"time"

	"fmt"

	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/server/mcstore/mcstoreapi"
)

type downloader struct {
	projectDB  ProjectDB
	dir        *Directory
	file       *File
	serverFile *mcstoreapi.ServerFile
	c          *ClientAPI
}

func newDownloader(projectDB ProjectDB, clientAPI *ClientAPI) *downloader {
	return &downloader{
		projectDB: projectDB,
		c:         clientAPI,
	}
}

func (d *downloader) downloadFile(path string) error {
	if err := os.MkdirAll(filepath.Dir(path), 0770); err != err {
		return err
	}
	project := d.projectDB.Project()
	var err error
	if d.serverFile, err = d.c.serverAPI.GetFileForPath(project.ProjectID, pathFromProject(path, project.Name)); err != nil {
		fmt.Println("getFileForPath returned error", err)
		return err
	}
	if d.dir, err = d.projectDB.FindDirectory(filepath.Dir(path)); err != nil {
		d.c.createDirectory(d.projectDB, filepath.Dir(path))
		d.dir, _ = d.projectDB.FindDirectory(filepath.Dir(path))
	}

	d.file, _ = d.projectDB.FindFile(filepath.Base(path), d.dir.ID)

	if finfo, err := os.Stat(path); os.IsNotExist(err) {
		return d.downloadNewFile(path)
	} else {
		return d.downloadExistingFile(finfo, path)
	}
}

func (d *downloader) downloadNewFile(path string) error {
	project := d.projectDB.Project()
	if err := d.c.serverAPI.DownloadFile(project.ProjectID, d.serverFile.ID, path); err != nil {
		fmt.Println("serverAPI.DownloadFile error", err)
		return err
	}
	finfo, _ := os.Stat(path)
	if d.file == nil {
		// File not in database

		newFile := File{
			FileID:     d.serverFile.ID,
			Name:       filepath.Base(path),
			Checksum:   d.serverFile.Checksum,
			Size:       finfo.Size(),
			MTime:      finfo.ModTime(),
			LastUpload: time.Now(),
			Directory:  d.dir.ID,
		}
		d.projectDB.InsertFile(&newFile)
	} else {
		// Existing file but not on users system
		d.file.MTime = finfo.ModTime()
		d.file.LastUpload = time.Now()
		d.file.FileID = d.serverFile.ID
		d.file.Size = finfo.Size()
		d.file.Checksum = d.serverFile.Checksum
		d.projectDB.UpdateFile(d.file)
	}
	return nil
}

func (d *downloader) downloadExistingFile(finfo os.FileInfo, path string) error {
	switch {
	case d.file == nil:
		// There is an existing file that isn't in database. Don't overwrite.
		fmt.Println("downloadExistingFile ErrFileNotUploaded")
		return ErrFileNotUploaded
	case finfo.ModTime().Unix() > d.file.MTime.Unix():
		// Existing file with updates that haven't been uploaded. Don't overwrite.
		fmt.Println("downloadExistingFile ErrFileVersionNotUploaded")
		return ErrFileVersionNotUploaded
	case d.file.Checksum == d.serverFile.Checksum:
		// Latest file already downloaded
		fmt.Println("downloadExistingFile Checksums are equal")
		return nil
	default:
		return d.downloadNewFile(path)
	}
}

func pathFromProject(path, projectName string) string {
	index := strings.Index(path, projectName)
	return path[index:len(path)]
}

type fentry struct {
	Type     string
	ID       string
	Path     string
	Size     int64
	Checksum string
}

type projectDownloader struct {
	downloader *downloader
	files      []fentry
}

func newProjectDownloader(projectDB ProjectDB, clientAPI *ClientAPI) *projectDownloader {
	return &projectDownloader{
		downloader: newDownloader(projectDB, clientAPI),
		files:      []fentry{},
	}
}

func (d *projectDownloader) downloadProject() error {
	project := d.downloader.projectDB.Project()

	if dir, err := d.downloader.c.getProjectDirList(project.ProjectID, ""); err == nil {
		d.files = append(d.files, toFentry(dir))
		d.loadDirRecurse(project.ProjectID, dir)
	}

	// Project Path contains the name of the project. The path for each entry
	// start with the project name. So we remove the project name from the
	// project path since the entry path will contain it.
	// eg, project path: /home/me/projects/PROJECT_NAME
	// individual entry paths: PROJECT_NAME/myfile.txt
	// so projectDir removes PROJECT_NAME, since joining with entry.Path will
	// put the PROJECT_NAME back in to the path.
	projectDir := filepath.Dir(project.Path)
	for _, e := range d.files {
		if e.Type == "file" {
			fmt.Println("Downloading file to:", filepath.Join(projectDir, e.Path))
			d.downloader.downloadFile(filepath.Join(projectDir, e.Path))
		} else if e.Type == "directory" {
			fmt.Println("Creating directory:", filepath.Join(projectDir, e.Path))
			d.createDir(filepath.Join(projectDir, e.Path), e.ID)
		}
	}
	return nil
}

func (d *projectDownloader) loadDirRecurse(projectID string, dentry *mcstoreapi.ServerDir) {
	if dentry.Type == "directory" {
		if dir, err := d.downloader.c.getProjectDirList(projectID, dentry.ID); err == nil {
			for _, entry := range dir.Children {
				d.files = append(d.files, toFentry(&entry))
				d.loadDirRecurse(projectID, &entry)
			}
		}
	}
}

func (d *projectDownloader) createDir(path, dirID string) {
	os.MkdirAll(path, 0770)
	if _, err := d.downloader.projectDB.FindDirectory(path); err == app.ErrNotFound {
		dir := &Directory{
			DirectoryID: dirID,
			Path:        path,
		}
		d.downloader.projectDB.InsertDirectory(dir)
	}
}

func toFentry(dentry *mcstoreapi.ServerDir) fentry {
	entry := fentry{
		Type:     dentry.Type,
		ID:       dentry.ID,
		Path:     dentry.Path,
		Size:     dentry.Size,
		Checksum: dentry.Checksum,
	}
	return entry
}
<commit_msg>Convert windows \ to /.<commit_after>package mc

import (
	"os"
	"path/filepath"
	"strings"

	"time"

	"fmt"

	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/server/mcstore/mcstoreapi"
)

type downloader struct {
	projectDB  ProjectDB
	dir        *Directory
	file       *File
	serverFile *mcstoreapi.ServerFile
	c          *ClientAPI
}

func newDownloader(projectDB ProjectDB, clientAPI *ClientAPI) *downloader {
	return &downloader{
		projectDB: projectDB,
		c:         clientAPI,
	}
}

func (d *downloader) downloadFile(path string) error {
	if err := os.MkdirAll(filepath.Dir(path), 0770); err != err {
		return err
	}
	project := d.projectDB.Project()
	var err error
	if d.serverFile, err = d.c.serverAPI.GetFileForPath(project.ProjectID, pathFromProject(path, project.Name)); err != nil {
		fmt.Println("getFileForPath returned error", err)
		return err
	}
	if d.dir, err = d.projectDB.FindDirectory(filepath.Dir(path)); err != nil {
		d.c.createDirectory(d.projectDB, filepath.Dir(path))
		d.dir, _ = d.projectDB.FindDirectory(filepath.Dir(path))
	}

	d.file, _ = d.projectDB.FindFile(filepath.Base(path), d.dir.ID)

	if finfo, err := os.Stat(path); os.IsNotExist(err) {
		return d.downloadNewFile(path)
	} else {
		return d.downloadExistingFile(finfo, path)
	}
}

func (d *downloader) downloadNewFile(path string) error {
	project := d.projectDB.Project()
	if err := d.c.serverAPI.DownloadFile(project.ProjectID, d.serverFile.ID, path); err != nil {
		fmt.Println("serverAPI.DownloadFile error", err)
		return err
	}
	finfo, _ := os.Stat(path)
	if d.file == nil {
		// File not in database

		newFile := File{
			FileID:     d.serverFile.ID,
			Name:       filepath.Base(path),
			Checksum:   d.serverFile.Checksum,
			Size:       finfo.Size(),
			MTime:      finfo.ModTime(),
			LastUpload: time.Now(),
			Directory:  d.dir.ID,
		}
		d.projectDB.InsertFile(&newFile)
	} else {
		// Existing file but not on users system
		d.file.MTime = finfo.ModTime()
		d.file.LastUpload = time.Now()
		d.file.FileID = d.serverFile.ID
		d.file.Size = finfo.Size()
		d.file.Checksum = d.serverFile.Checksum
		d.projectDB.UpdateFile(d.file)
	}
	return nil
}

func (d *downloader) downloadExistingFile(finfo os.FileInfo, path string) error {
	switch {
	case d.file == nil:
		// There is an existing file that isn't in database. Don't overwrite.
		fmt.Println("downloadExistingFile ErrFileNotUploaded")
		return ErrFileNotUploaded
	case finfo.ModTime().Unix() > d.file.MTime.Unix():
		// Existing file with updates that haven't been uploaded. Don't overwrite.
		fmt.Println("downloadExistingFile ErrFileVersionNotUploaded")
		return ErrFileVersionNotUploaded
	case d.file.Checksum == d.serverFile.Checksum:
		// Latest file already downloaded
		fmt.Println("downloadExistingFile Checksums are equal")
		return nil
	default:
		return d.downloadNewFile(path)
	}
}

func pathFromProject(path, projectName string) string {
	index := strings.Index(path, projectName)
	return strings.Replace(path[index:len(path)], "\\", "/", -1)
}

type fentry struct {
	Type     string
	ID       string
	Path     string
	Size     int64
	Checksum string
}

type projectDownloader struct {
	downloader *downloader
	files      []fentry
}

func newProjectDownloader(projectDB ProjectDB, clientAPI *ClientAPI) *projectDownloader {
	return &projectDownloader{
		downloader: newDownloader(projectDB, clientAPI),
		files:      []fentry{},
	}
}

func (d *projectDownloader) downloadProject() error {
	project := d.downloader.projectDB.Project()

	if dir, err := d.downloader.c.getProjectDirList(project.ProjectID, ""); err == nil {
		d.files = append(d.files, toFentry(dir))
		d.loadDirRecurse(project.ProjectID, dir)
	}

	// Project Path contains the name of the project. The path for each entry
	// start with the project name. So we remove the project name from the
	// project path since the entry path will contain it.
	// eg, project path: /home/me/projects/PROJECT_NAME
	// individual entry paths: PROJECT_NAME/myfile.txt
	// so projectDir removes PROJECT_NAME, since joining with entry.Path will
	// put the PROJECT_NAME back in to the path.
	projectDir := filepath.Dir(project.Path)
	for _, e := range d.files {
		if e.Type == "file" {
			fmt.Println("Downloading file to:", filepath.Join(projectDir, e.Path))
			d.downloader.downloadFile(filepath.Join(projectDir, e.Path))
		} else if e.Type == "directory" {
			fmt.Println("Creating directory:", filepath.Join(projectDir, e.Path))
			d.createDir(filepath.Join(projectDir, e.Path), e.ID)
		}
	}
	return nil
}

func (d *projectDownloader) loadDirRecurse(projectID string, dentry *mcstoreapi.ServerDir) {
	if dentry.Type == "directory" {
		if dir, err := d.downloader.c.getProjectDirList(projectID, dentry.ID); err == nil {
			for _, entry := range dir.Children {
				d.files = append(d.files, toFentry(&entry))
				d.loadDirRecurse(projectID, &entry)
			}
		}
	}
}

func (d *projectDownloader) createDir(path, dirID string) {
	os.MkdirAll(path, 0770)
	if _, err := d.downloader.projectDB.FindDirectory(path); err == app.ErrNotFound {
		dir := &Directory{
			DirectoryID: dirID,
			Path:        path,
		}
		d.downloader.projectDB.InsertDirectory(dir)
	}
}

func toFentry(dentry *mcstoreapi.ServerDir) fentry {
	entry := fentry{
		Type:     dentry.Type,
		ID:       dentry.ID,
		Path:     dentry.Path,
		Size:     dentry.Size,
		Checksum: dentry.Checksum,
	}
	return entry
}
<|endoftext|>


<commit_before>package main

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

	// TODO(mc) Send to airbrake
}
<commit_msg>Remove error.go<commit_after><|endoftext|>


<commit_before>package vsock

import (
	"errors"
	"os"
	"reflect"
	"testing"

	"golang.org/x/sys/unix"
)

func Test_listenStreamLinuxHandleError(t *testing.T) {
	var closed bool
	lfd := &testFD{
		// Track when fd.Close is called.
		close: func() error {
			closed = true
			return nil
		},
		// Always return an error on bind.
		bind: func(sa unix.Sockaddr) error {
			return errors.New("error during bind")
		},
	}

	if _, err := listenStreamLinuxHandleError(lfd, 0, 0); err == nil {
		t.Fatal("expected an error, but none occurred")
	}

	if want, got := true, closed; want != got {
		t.Fatalf("unexpected socket close value:\n- want: %v\n-  got: %v",
			want, got)
	}
}

func Test_listenStreamLinuxPortZero(t *testing.T) {
	const (
		cid  uint32 = ContextIDHost
		port uint32 = 0
	)

	lsa := &unix.SockaddrVM{
		CID: cid,
		// Expect 0 to be turned into "any port".
		Port: unix.VMADDR_PORT_ANY,
	}

	bindFn := func(sa unix.Sockaddr) error {
		if want, got := lsa, sa; !reflect.DeepEqual(want, got) {
			t.Fatalf("unexpected bind sockaddr:\n- want: %#v\n-  got: %#v",
				want, got)
		}

		return nil
	}

	lfd := &testFD{
		bind:        bindFn,
		listen:      func(n int) error { return nil },
		getsockname: func() (unix.Sockaddr, error) { return lsa, nil },
	}

	if _, err := listenStreamLinux(lfd, cid, port); err != nil {
		t.Fatalf("failed to listen: %v", err)
	}
}

func Test_listenStreamLinuxFull(t *testing.T) {
	const (
		cid  uint32 = ContextIDHost
		port uint32 = 1024
	)

	lsa := &unix.SockaddrVM{
		CID:  cid,
		Port: port,
	}

	bindFn := func(sa unix.Sockaddr) error {
		if want, got := lsa, sa; !reflect.DeepEqual(want, got) {
			t.Fatalf("unexpected bind sockaddr:\n- want: %#v\n-  got: %#v",
				want, got)
		}

		return nil
	}

	listenFn := func(n int) error {
		if want, got := listenBacklog, n; want != got {
			t.Fatalf("unexpected listen backlog:\n- want: %d\n-  got: %d",
				want, got)
		}

		return nil
	}

	lfd := &testFD{
		bind:   bindFn,
		listen: listenFn,
		getsockname: func() (unix.Sockaddr, error) {
			return lsa, nil
		},
	}

	nl, err := listenStreamLinux(lfd, cid, port)
	if err != nil {
		t.Fatalf("failed to listen: %v", err)
	}

	l := nl.(*listener)

	if want, got := cid, l.addr.ContextID; want != got {
		t.Fatalf("unexpected listener context ID:\n- want: %d\n-  got: %d",
			want, got)
	}
	if want, got := port, l.addr.Port; want != got {
		t.Fatalf("unexpected listener context ID:\n- want: %d\n-  got: %d",
			want, got)
	}
}

func Test_listenerAccept(t *testing.T) {
	const (
		connFD uintptr = 10

		cid  uint32 = 3
		port uint32 = 1024
	)

	accept4Fn := func(flags int) (fd, unix.Sockaddr, error) {
		if want, got := 0, flags; want != got {
			t.Fatalf("unexpected accept4 flags:\n- want: %d\n-  got: %d",
				want, got)
		}

		acceptFD := &testFD{
			newFile: func(name string) *os.File {
				return os.NewFile(connFD, name)
			},
		}

		acceptSA := &unix.SockaddrVM{
			CID:  cid,
			Port: port,
		}

		return acceptFD, acceptSA, nil
	}

	localAddr := &Addr{
		ContextID: ContextIDHost,
		Port:      port,
	}

	l := &listener{
		fd: &testFD{
			accept4: accept4Fn,
		},
		addr: localAddr,
	}

	nc, err := l.Accept()
	if err != nil {
		t.Fatalf("failed to accept: %v", err)
	}

	c := nc.(*conn)

	if want, got := localAddr, c.LocalAddr(); !reflect.DeepEqual(want, got) {
		t.Fatalf("unexpected conn local address:\n- want: %#v\n-  got: %#v",
			want, got)
	}

	remoteAddr := &Addr{
		ContextID: cid,
		Port:      port,
	}

	if want, got := remoteAddr, c.RemoteAddr(); !reflect.DeepEqual(want, got) {
		t.Fatalf("unexpected conn remote address:\n- want: %#v\n-  got: %#v",
			want, got)
	}

	if want, got := connFD, c.file.Fd(); want != got {
		t.Fatalf("unexpected conn file descriptor:\n- want: %d\n-  got: %d",
			want, got)
	}
}
<commit_msg>vsock: adjust listener test for nonblocking<commit_after>package vsock

import (
	"errors"
	"os"
	"reflect"
	"testing"

	"golang.org/x/sys/unix"
)

func Test_listenStreamLinuxHandleError(t *testing.T) {
	var closed bool
	lfd := &testFD{
		// Track when fd.Close is called.
		close: func() error {
			closed = true
			return nil
		},
		// Always return an error on bind.
		bind: func(sa unix.Sockaddr) error {
			return errors.New("error during bind")
		},
	}

	if _, err := listenStreamLinuxHandleError(lfd, 0, 0); err == nil {
		t.Fatal("expected an error, but none occurred")
	}

	if want, got := true, closed; want != got {
		t.Fatalf("unexpected socket close value:\n- want: %v\n-  got: %v",
			want, got)
	}
}

func Test_listenStreamLinuxPortZero(t *testing.T) {
	const (
		cid  uint32 = ContextIDHost
		port uint32 = 0
	)

	lsa := &unix.SockaddrVM{
		CID: cid,
		// Expect 0 to be turned into "any port".
		Port: unix.VMADDR_PORT_ANY,
	}

	bindFn := func(sa unix.Sockaddr) error {
		if want, got := lsa, sa; !reflect.DeepEqual(want, got) {
			t.Fatalf("unexpected bind sockaddr:\n- want: %#v\n-  got: %#v",
				want, got)
		}

		return nil
	}

	lfd := &testFD{
		bind:        bindFn,
		listen:      func(n int) error { return nil },
		getsockname: func() (unix.Sockaddr, error) { return lsa, nil },
	}

	if _, err := listenStreamLinux(lfd, cid, port); err != nil {
		t.Fatalf("failed to listen: %v", err)
	}
}

func Test_listenStreamLinuxFull(t *testing.T) {
	const (
		cid  uint32 = ContextIDHost
		port uint32 = 1024
	)

	lsa := &unix.SockaddrVM{
		CID:  cid,
		Port: port,
	}

	lfd := &testFD{
		bind: func(sa unix.Sockaddr) error {
			if want, got := lsa, sa; !reflect.DeepEqual(want, got) {
				t.Fatalf("unexpected bind sockaddr:\n- want: %#v\n-  got: %#v",
					want, got)
			}

			return nil
		},
		listen: func(n int) error {
			if want, got := listenBacklog, n; want != got {
				t.Fatalf("unexpected listen backlog:\n- want: %d\n-  got: %d",
					want, got)
			}

			return nil
		},
		getsockname: func() (unix.Sockaddr, error) {
			return lsa, nil
		},
		setNonblock: func(nonblocking bool) error {
			if want, got := true, nonblocking; !reflect.DeepEqual(want, got) {
				t.Fatalf("unexpected set nonblocking value:\n- want: %#v\n-  got: %#v",
					want, got)
			}

			return nil
		},
	}

	nl, err := listenStreamLinux(lfd, cid, port)
	if err != nil {
		t.Fatalf("failed to listen: %v", err)
	}

	l := nl.(*listener)

	if want, got := cid, l.addr.ContextID; want != got {
		t.Fatalf("unexpected listener context ID:\n- want: %d\n-  got: %d",
			want, got)
	}
	if want, got := port, l.addr.Port; want != got {
		t.Fatalf("unexpected listener context ID:\n- want: %d\n-  got: %d",
			want, got)
	}
}

func Test_listenerAccept(t *testing.T) {
	const (
		connFD uintptr = 10

		cid  uint32 = 3
		port uint32 = 1024
	)

	accept4Fn := func(flags int) (fd, unix.Sockaddr, error) {
		if want, got := 0, flags; want != got {
			t.Fatalf("unexpected accept4 flags:\n- want: %d\n-  got: %d",
				want, got)
		}

		acceptFD := &testFD{
			newFile: func(name string) *os.File {
				return os.NewFile(connFD, name)
			},
			setNonblock: func(nonblocking bool) error {
				if want, got := true, nonblocking; !reflect.DeepEqual(want, got) {
					t.Fatalf("unexpected set nonblocking value:\n- want: %#v\n-  got: %#v",
						want, got)
				}
				return nil
			},
		}

		acceptSA := &unix.SockaddrVM{
			CID:  cid,
			Port: port,
		}

		return acceptFD, acceptSA, nil
	}

	localAddr := &Addr{
		ContextID: ContextIDHost,
		Port:      port,
	}

	l := &listener{
		fd: &testFD{
			accept4: accept4Fn,
		},
		addr: localAddr,
	}

	nc, err := l.Accept()
	if err != nil {
		t.Fatalf("failed to accept: %v", err)
	}

	c := nc.(*conn)

	if want, got := localAddr, c.LocalAddr(); !reflect.DeepEqual(want, got) {
		t.Fatalf("unexpected conn local address:\n- want: %#v\n-  got: %#v",
			want, got)
	}

	remoteAddr := &Addr{
		ContextID: cid,
		Port:      port,
	}

	if want, got := remoteAddr, c.RemoteAddr(); !reflect.DeepEqual(want, got) {
		t.Fatalf("unexpected conn remote address:\n- want: %#v\n-  got: %#v",
			want, got)
	}

	if want, got := connFD, c.file.Fd(); want != got {
		t.Fatalf("unexpected conn file descriptor:\n- want: %d\n-  got: %d",
			want, got)
	}
}
<|endoftext|>