<commit_before>package sourcegraph

import (
	"fmt"

	"github.com/sourcegraph/vcsstore/vcsclient"

	"sourcegraph.com/sourcegraph/go-sourcegraph/router"
)

// RepositoryTreeService communicates with the Sourcegraph API endpoints that
// fetch file and directory entries in repositories.
type RepositoryTreeService interface {
	Get(entry TreeEntrySpec, opt *RepositoryTreeGetOptions) (*TreeEntry, Response, error)
}

type repositoryTreeService struct {
	client *Client
}

var _ RepositoryTreeService = &repositoryTreeService{}

type TreeEntrySpec struct {
	RepoRev RepoRevSpec
	Path    string
}

func (s *TreeEntrySpec) RouteVars() map[string]string {
	m := s.RepoRev.RouteVars()
	m["Path"] = s.Path
	return m
}

func (s TreeEntrySpec) String() string {
	return fmt.Sprintf("%v: %s (rev %q)", s.RepoRev, s.Path, s.RepoRev.Rev)
}

// TreeEntry is a file or directory in a repository, with additional feedback
// from the formatting operation (if Formatted is true in the options).
type TreeEntry struct {
	*vcsclient.TreeEntry

	ContentsString string

	// FormatResult is only set if this TreeEntry is a file.
	FormatResult *FormatResult `json:",omitempty"`

	// EntryDefinitions is a list of defined defs for each entry in this
	// directory. It is only populated if DirEntryDefinitions is true.
	EntryDefinitions map[string]interface{}
}

// FormatResult contains information about and warnings from the formatting
// operation (if Formatted is true in the options).
type FormatResult struct {
	// TooManyRefs indicates that the file being formatted exceeded the maximum
	// number of refs that are linked. Only the first NumRefs refs are linked.
	TooManyRefs bool `json:",omitempty"`

	// NumRefs is the number of refs that were linked in this file. If the total
	// number of refs in the file exceeds the (server-defined) limit, NumRefs is
	// capped at the limit.
	NumRefs int

	// The line in the file that the formatted section starts at
	StartLine int

	// The line that the formatted section ends at
	EndLine int
}

// RepositoryTreeGetOptions specifies options for (RepositoryTreeService).Get.
type RepositoryTreeGetOptions struct {
	// Formatted is whether the specified entry, if it's a file, should have its
	// contents code-formatted.
	Formatted bool

	// DirEntryDefinitions is whether the specified entry, if it's a directory,
	// should include a list of defined defs for each of its entries (in
	// EntryDefinitions). For example, if the specified entry has a file "a" and
	// a dir "b/", the result would include a list of defs defined in "a" and
	// in any file underneath "b/". Not all defs defined in the entries are
	// returned; only the top few are.
	DirEntryDefinitions bool `url:",omitempty"`

	ContentsAsString bool `url:",omitempty"`

	vcsclient.GetFileOptions
}

func (s *repositoryTreeService) Get(entry TreeEntrySpec, opt *RepositoryTreeGetOptions) (*TreeEntry, Response, error) {
	url, err := s.client.url(router.RepositoryTreeEntry, entry.RouteVars(), opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", url.String(), nil)
	if err != nil {
		return nil, nil, err
	}

	var entry_ *TreeEntry
	resp, err := s.client.Do(req, &entry_)
	if err != nil {
		return nil, resp, err
	}

	return entry_, resp, nil
}

type MockRepositoryTreeService struct {
	Get_ func(entry TreeEntrySpec, opt *RepositoryTreeGetOptions) (*TreeEntry, Response, error)
}

var _ RepositoryTreeService = MockRepositoryTreeService{}

func (s MockRepositoryTreeService) Get(entry TreeEntrySpec, opt *RepositoryTreeGetOptions) (*TreeEntry, Response, error) {
	if s.Get_ == nil {
		return nil, &HTTPResponse{}, nil
	}
	return s.Get_(entry, opt)
}
<commit_msg>remove fields that are now in vcsclient.FileWithRange<commit_after>package sourcegraph

import (
	"fmt"

	"github.com/sourcegraph/vcsstore/vcsclient"

	"sourcegraph.com/sourcegraph/go-sourcegraph/router"
)

// RepositoryTreeService communicates with the Sourcegraph API endpoints that
// fetch file and directory entries in repositories.
type RepositoryTreeService interface {
	Get(entry TreeEntrySpec, opt *RepositoryTreeGetOptions) (*TreeEntry, Response, error)
}

type repositoryTreeService struct {
	client *Client
}

var _ RepositoryTreeService = &repositoryTreeService{}

type TreeEntrySpec struct {
	RepoRev RepoRevSpec
	Path    string
}

func (s *TreeEntrySpec) RouteVars() map[string]string {
	m := s.RepoRev.RouteVars()
	m["Path"] = s.Path
	return m
}

func (s TreeEntrySpec) String() string {
	return fmt.Sprintf("%v: %s (rev %q)", s.RepoRev, s.Path, s.RepoRev.Rev)
}

// TreeEntry is a file or directory in a repository, with additional feedback
// from the formatting operation (if Formatted is true in the options).
type TreeEntry struct {
	*vcsclient.TreeEntry

	ContentsString string

	// FormatResult is only set if this TreeEntry is a file.
	FormatResult *FormatResult `json:",omitempty"`

	// EntryDefinitions is a list of defined defs for each entry in this
	// directory. It is only populated if DirEntryDefinitions is true.
	EntryDefinitions map[string]interface{}
}

// FormatResult contains information about and warnings from the formatting
// operation (if Formatted is true in the options).
type FormatResult struct {
	// TooManyRefs indicates that the file being formatted exceeded the maximum
	// number of refs that are linked. Only the first NumRefs refs are linked.
	TooManyRefs bool `json:",omitempty"`

	// NumRefs is the number of refs that were linked in this file. If the total
	// number of refs in the file exceeds the (server-defined) limit, NumRefs is
	// capped at the limit.
	NumRefs int
}

// RepositoryTreeGetOptions specifies options for (RepositoryTreeService).Get.
type RepositoryTreeGetOptions struct {
	// Formatted is whether the specified entry, if it's a file, should have its
	// contents code-formatted.
	Formatted bool

	// DirEntryDefinitions is whether the specified entry, if it's a directory,
	// should include a list of defined defs for each of its entries (in
	// EntryDefinitions). For example, if the specified entry has a file "a" and
	// a dir "b/", the result would include a list of defs defined in "a" and
	// in any file underneath "b/". Not all defs defined in the entries are
	// returned; only the top few are.
	DirEntryDefinitions bool `url:",omitempty"`

	ContentsAsString bool `url:",omitempty"`

	vcsclient.GetFileOptions
}

func (s *repositoryTreeService) Get(entry TreeEntrySpec, opt *RepositoryTreeGetOptions) (*TreeEntry, Response, error) {
	url, err := s.client.url(router.RepositoryTreeEntry, entry.RouteVars(), opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", url.String(), nil)
	if err != nil {
		return nil, nil, err
	}

	var entry_ *TreeEntry
	resp, err := s.client.Do(req, &entry_)
	if err != nil {
		return nil, resp, err
	}

	return entry_, resp, nil
}

type MockRepositoryTreeService struct {
	Get_ func(entry TreeEntrySpec, opt *RepositoryTreeGetOptions) (*TreeEntry, Response, error)
}

var _ RepositoryTreeService = MockRepositoryTreeService{}

func (s MockRepositoryTreeService) Get(entry TreeEntrySpec, opt *RepositoryTreeGetOptions) (*TreeEntry, Response, error) {
	if s.Get_ == nil {
		return nil, &HTTPResponse{}, nil
	}
	return s.Get_(entry, opt)
}
<|endoftext|>


<commit_before>// Copyright 2017 the u-root 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 main

import (
	"bytes"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/u-root/u-root/pkg/testutil"
)

type WifiTestCase struct {
	name      string
	args      []string
	out       string
	exp_error bool
}

var testcases = []WifiTestCase{
	{
		name:      "No Flags, No Args",
		args:      nil,
		out:       "Usage",
		exp_error: true,
	},
	{
		name:      "Flags, No Args",
		args:      []string{"-i=123"},
		out:       "Usage",
		exp_error: true,
	},
}

var ERROR_MSG_FORMAT = "\nEXPECTED:\n%s\n\nACTUAL:\n%s\n"

func errorExists(err error) bool {
	return err != nil
}

func craftPrintMsg(err_exists bool, out string) string {
	var msg bytes.Buffer

	if err_exists {
		msg.WriteString("Error Status: exists\n")
	} else {
		msg.WriteString("Error Status: not exists\n")
	}
	msg.WriteString("Output:\n")
	msg.WriteString(out)
	return msg.String()
}

func TestWifi(t *testing.T) {
	// Set up
	tmpDir, execPath := testutil.CompileInTempDir(t)
	defer os.RemoveAll(tmpDir)

	// Tests
	for _, test := range testcases {
		t.Logf("TEST %v", test.name)
		c := exec.Command(execPath, test.args...)
		out, err := c.CombinedOutput()
		if (test.exp_error != errorExists(err)) || !strings.Contains(string(out), test.out) {
			expectMsg := craftPrintMsg(test.exp_error, test.out)
			actualMsg := craftPrintMsg(errorExists(err), string(out))
			t.Errorf(ERROR_MSG_FORMAT, expectMsg, actualMsg)
		}
	}
}
<commit_msg>Fixed wifi_test.go to agree with Go convention<commit_after>// Copyright 2017 the u-root 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 main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/u-root/u-root/pkg/testutil"
)

type WifiTestCase struct {
	name      string
	args      []string
	out       string
	errExists bool
}

var testcases = []WifiTestCase{
	{
		name:      "No Flags, No Args",
		args:      nil,
		out:       "Usage",
		errExists: true,
	},
	{
		name:      "Flags, No Args",
		args:      []string{"-i=123"},
		out:       "Usage",
		errExists: true,
	},
}

func errorExists(err error) bool {
	return err != nil
}

func craftPrintMsg(errExists bool, out string) string {
	var msg bytes.Buffer

	if errExists {
		msg.WriteString("Error Status: exists\n")
	} else {
		msg.WriteString("Error Status: not exists\n")
	}
	msg.WriteString("Output:\n")
	msg.WriteString(out)
	return msg.String()
}

func TestWifi(t *testing.T) {
	// Set up
	tmpDir, execPath := testutil.CompileInTempDir(t)
	defer os.RemoveAll(tmpDir)

	// Tests
	for _, test := range testcases {
		t.Logf("TEST %v", test.name)
		c := exec.Command(execPath, test.args...)
		out, err := c.CombinedOutput()
		if (test.errExists != errorExists(err)) || !strings.Contains(string(out), test.out) {
			expectMsg := craftPrintMsg(test.errExists, test.out)
			actualMsg := craftPrintMsg(errorExists(err), string(out))
			execStatement := fmt.Sprintf("exec(wifi %s)", strings.Trim(fmt.Sprint(test.args), "[]"))
			t.Errorf("%s\ngot:\n%s\n\nwant:\n%s", execStatement, actualMsg, expectMsg)
		}
	}
}
<|endoftext|>


<commit_before>package snapshot

import (
	"encoding/binary"
	"errors"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"math/rand"
	"sync/atomic"
	"time"

	"github.com/kopia/kopia/fs"
	"github.com/kopia/kopia/internal/dir"
	"github.com/kopia/kopia/internal/hashcache"
	"github.com/kopia/kopia/repo"
)

func hashEntryMetadata(w io.Writer, e *fs.EntryMetadata) {
	binary.Write(w, binary.LittleEndian, e.Name)
	binary.Write(w, binary.LittleEndian, e.ModTime.UnixNano())
	binary.Write(w, binary.LittleEndian, e.FileMode())
	binary.Write(w, binary.LittleEndian, e.FileSize)
	binary.Write(w, binary.LittleEndian, e.UserID)
	binary.Write(w, binary.LittleEndian, e.GroupID)
}

func metadataHash(e *fs.EntryMetadata) uint64 {
	h := fnv.New64a()
	hashEntryMetadata(h, e)
	return h.Sum64()
}

var errCancelled = errors.New("cancelled")

// Uploader supports efficient uploading files and directories to repository.
type Uploader struct {
	Progress               UploadProgress
	Files                  FilesPolicy
	MaxUploadBytes         int64
	IgnoreFileErrors       bool
	ForceHashingPercentage int

	uploadBuf   []byte
	repo        *repo.Repository
	cacheWriter hashcache.Writer
	cacheReader hashcache.Reader

	stats     Stats
	cancelled int32
}

// IsCancelled returns true if the upload is cancelled.
func (u *Uploader) IsCancelled() bool {
	return u.cancelReason() != ""
}

func (u *Uploader) cancelReason() string {
	if c := atomic.LoadInt32(&u.cancelled) != 0; c {
		return "cancelled"
	}

	if mub := u.MaxUploadBytes; mub > 0 && u.repo.Stats().WrittenBytes > mub {
		return "limit reached"
	}

	return ""
}

func (u *Uploader) uploadFileInternal(f fs.File, relativePath string, forceStored bool) (*dir.Entry, uint64, error) {
	u.Progress.Started(relativePath, f.Metadata().FileSize)

	file, err := f.Open()
	if err != nil {
		return nil, 0, fmt.Errorf("unable to open file: %v", err)
	}
	defer file.Close()

	writer := u.repo.NewWriter(repo.WriterOptions{
		Description: "FILE:" + f.Metadata().Name,
	})
	defer writer.Close()

	written, err := u.copyWithProgress(relativePath, writer, file, 0, f.Metadata().FileSize)
	if err != nil {
		u.Progress.Finished(relativePath, f.Metadata().FileSize, err)
		return nil, 0, err
	}

	e2, err := file.EntryMetadata()
	if err != nil {
		u.Progress.Finished(relativePath, f.Metadata().FileSize, err)
		return nil, 0, err
	}

	r, err := writer.Result(forceStored)
	if err != nil {
		u.Progress.Finished(relativePath, f.Metadata().FileSize, err)
		return nil, 0, err
	}

	de := newDirEntry(e2, r)
	de.FileSize = written

	u.Progress.Finished(relativePath, f.Metadata().FileSize, nil)

	return de, metadataHash(&de.EntryMetadata), nil
}

func (u *Uploader) copyWithProgress(path string, dst io.Writer, src io.Reader, completed int64, length int64) (int64, error) {
	if u.uploadBuf == nil {
		u.uploadBuf = make([]byte, 128*1024) // 128 KB buffer
	}

	var written int64

	for {
		if u.IsCancelled() {
			return 0, errCancelled
		}

		readBytes, readErr := src.Read(u.uploadBuf)
		if readBytes > 0 {
			wroteBytes, writeErr := dst.Write(u.uploadBuf[0:readBytes])
			if wroteBytes > 0 {
				written += int64(wroteBytes)
				completed += int64(wroteBytes)
				if length < completed {
					length = completed
				}
				u.Progress.Progress(path, completed, length)
			}
			if writeErr != nil {
				return written, writeErr
			}
			if readBytes != wroteBytes {
				return written, io.ErrShortWrite
			}
		}

		if readErr != nil {
			if readErr == io.EOF {
				break
			}

			return written, readErr
		}
	}

	return written, nil
}

func newDirEntry(md *fs.EntryMetadata, oid repo.ObjectID) *dir.Entry {
	return &dir.Entry{
		EntryMetadata: *md,
		ObjectID:      oid,
	}
}

// uploadFile uploads the specified File to the repository.
func (u *Uploader) uploadFile(file fs.File) (repo.ObjectID, error) {
	e, _, err := u.uploadFileInternal(file, file.Metadata().Name, true)
	if err != nil {
		return repo.NullObjectID, err
	}
	return e.ObjectID, nil
}

// uploadDir uploads the specified Directory to the repository.
// An optional ID of a hash-cache object may be provided, in which case the Uploader will use its
// contents to avoid hashing
func (u *Uploader) uploadDir(dir fs.Directory) (repo.ObjectID, repo.ObjectID, error) {
	var err error

	if err := u.repo.BeginPacking(); err != nil {
		return repo.NullObjectID, repo.NullObjectID, err
	}

	mw := u.repo.NewWriter(repo.WriterOptions{
		Description:     "HASHCACHE:" + dir.Metadata().Name,
		BlockNamePrefix: "H",
		PackGroup:       "HC",
	})
	defer mw.Close()
	u.cacheWriter = hashcache.NewWriter(mw)
	oid, err := uploadDirInternal(u, dir, ".", true)
	if u.IsCancelled() {
		if err := u.cacheReader.CopyTo(u.cacheWriter); err != nil {
			return repo.NullObjectID, repo.NullObjectID, err
		}
	}
	u.cacheWriter.Finalize()
	u.cacheWriter = nil

	if err != nil {
		return repo.NullObjectID, repo.NullObjectID, err
	}

	hcid, err := mw.Result(true)
	if err := u.repo.FinishPacking(); err != nil {
		return repo.NullObjectID, repo.NullObjectID, fmt.Errorf("can't finish packing: %v", err)
	}
	return oid, hcid, err
}

func uploadDirInternal(
	u *Uploader,
	directory fs.Directory,
	relativePath string,
	forceStored bool,
) (repo.ObjectID, error) {
	u.Progress.StartedDir(relativePath)
	defer u.Progress.FinishedDir(relativePath)

	u.stats.TotalDirectoryCount++

	entries, err := directory.Readdir()
	if err != nil {
		return repo.NullObjectID, err
	}

	writer := u.repo.NewWriter(repo.WriterOptions{
		Description: "DIR:" + relativePath,
		PackGroup:   "DIR",
	})

	dw := dir.NewWriter(writer)
	defer writer.Close()

	for _, entry := range entries {
		if u.IsCancelled() {
			break
		}
		e := entry.Metadata()
		entryRelativePath := relativePath + "/" + e.Name

		var de *dir.Entry

		var hash uint64

		switch entry := entry.(type) {
		case fs.Directory:
			oid, err := uploadDirInternal(u, entry, entryRelativePath, false)
			if err != nil {
				return repo.NullObjectID, err
			}
			de = newDirEntry(e, oid)

		case fs.Symlink:
			l, err := entry.Readlink()
			if err != nil {
				return repo.NullObjectID, err
			}

			de = newDirEntry(e, repo.InlineObjectID([]byte(l)))
			hash = metadataHash(e)

		case fs.File:
			// regular file
			// See if we had this name during previous pass.
			cachedEntry := u.maybeIgnoreHashCacheEntry(u.cacheReader.FindEntry(entryRelativePath))

			// ... and whether file metadata is identical to the previous one.
			computedHash := metadataHash(e)
			cacheMatches := (cachedEntry != nil) && cachedEntry.Hash == computedHash

			if cacheMatches {
				u.stats.CachedFiles++
				u.Progress.Cached(entryRelativePath, entry.Metadata().FileSize)
				// Avoid hashing by reusing previous object ID.
				de = newDirEntry(e, cachedEntry.ObjectID)
				hash = cachedEntry.Hash
			} else {
				u.stats.NonCachedFiles++
				de, hash, err = u.uploadFileInternal(entry, entryRelativePath, false)
				if err == errCancelled {
					continue
				}

				if err != nil {
					if u.IgnoreFileErrors {
						u.stats.ReadErrors++
						log.Printf("warning: unable to hash file %q: %s, ignoring", entryRelativePath, err)
						continue
					}
					return repo.NullObjectID, fmt.Errorf("unable to hash file: %s", err)
				}
			}

			u.stats.TotalFileCount++
			u.stats.TotalFileSize += de.FileSize

		default:
			return repo.NullObjectID, fmt.Errorf("file type %v not supported", entry.Metadata().Type)
		}

		if err := dw.WriteEntry(de); err != nil {
			return repo.NullObjectID, err
		}

		if de.Type != fs.EntryTypeDirectory && de.ObjectID.StorageBlock != "" {
			if err := u.cacheWriter.WriteEntry(hashcache.Entry{
				Name:     entryRelativePath,
				Hash:     hash,
				ObjectID: de.ObjectID,
			}); err != nil {
				return repo.NullObjectID, err
			}
		}
	}

	dw.Finalize()

	return writer.Result(forceStored)
}

func (u *Uploader) maybeIgnoreHashCacheEntry(e *hashcache.Entry) *hashcache.Entry {
	if rand.Intn(100) < u.ForceHashingPercentage {
		return nil
	}

	return e
}

// NewUploader creates new Uploader object for a given repository.
func NewUploader(r *repo.Repository) *Uploader {
	return &Uploader{
		repo:     r,
		Progress: &nullUploadProgress{},
	}
}

// Cancel requests cancellation of an upload that's in progress. Will typically result in an incomplete snapshot.
func (u *Uploader) Cancel() {
	atomic.StoreInt32(&u.cancelled, 1)
}

// Upload uploads contents of the specified filesystem entry (file or directory) to the repository and returns snapshot.Manifest with statistics.
// Old snapshot manifest, when provided can be used to speed up uploads by utilizing hash cache.
func (u *Uploader) Upload(
	source fs.Entry,
	sourceInfo *SourceInfo,
	old *Manifest,
) (*Manifest, error) {
	s := &Manifest{
		Source: *sourceInfo,
	}

	u.cacheReader = hashcache.Open(nil)
	u.stats = Stats{}
	if old != nil {
		if r, err := u.repo.Open(old.HashCacheID); err == nil {
			u.cacheReader = hashcache.Open(r)
		}
	}

	var err error

	s.StartTime = time.Now()

	switch entry := source.(type) {
	case fs.Directory:
		s.RootObjectID, s.HashCacheID, err = u.uploadDir(entry)

	case fs.File:
		s.RootObjectID, err = u.uploadFile(entry)

	default:
		return nil, fmt.Errorf("unsupported source: %v", s.Source)
	}
	if err != nil {
		return nil, err
	}

	s.IncompleteReason = u.cancelReason()
	s.EndTime = time.Now()
	s.Stats = u.stats
	s.Stats.Repository = u.repo.Status().Stats

	return s, nil
}
<commit_msg>simplified how upload code is structured + how hash cache is emitted<commit_after>package snapshot

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"math/rand"
	"sync/atomic"
	"time"

	"github.com/kopia/kopia/fs"
	"github.com/kopia/kopia/internal/dir"
	"github.com/kopia/kopia/internal/hashcache"
	"github.com/kopia/kopia/repo"
)

func hashEntryMetadata(w io.Writer, e *fs.EntryMetadata) {
	binary.Write(w, binary.LittleEndian, e.Name)
	binary.Write(w, binary.LittleEndian, e.ModTime.UnixNano())
	binary.Write(w, binary.LittleEndian, e.FileMode())
	binary.Write(w, binary.LittleEndian, e.FileSize)
	binary.Write(w, binary.LittleEndian, e.UserID)
	binary.Write(w, binary.LittleEndian, e.GroupID)
}

func metadataHash(e *fs.EntryMetadata) uint64 {
	h := fnv.New64a()
	hashEntryMetadata(h, e)
	return h.Sum64()
}

var errCancelled = errors.New("cancelled")

// Uploader supports efficient uploading files and directories to repository.
type Uploader struct {
	Progress               UploadProgress
	Files                  FilesPolicy
	MaxUploadBytes         int64
	IgnoreFileErrors       bool
	ForceHashingPercentage int

	uploadBuf   []byte
	repo        *repo.Repository
	cacheWriter hashcache.Writer
	cacheReader hashcache.Reader

	stats     Stats
	cancelled int32
}

// IsCancelled returns true if the upload is cancelled.
func (u *Uploader) IsCancelled() bool {
	return u.cancelReason() != ""
}

func (u *Uploader) cancelReason() string {
	if c := atomic.LoadInt32(&u.cancelled) != 0; c {
		return "cancelled"
	}

	if mub := u.MaxUploadBytes; mub > 0 && u.repo.Stats().WrittenBytes > mub {
		return "limit reached"
	}

	return ""
}

func (u *Uploader) uploadFileInternal(f fs.File, relativePath string, forceStored bool) (*dir.Entry, uint64, error) {
	u.Progress.Started(relativePath, f.Metadata().FileSize)

	file, err := f.Open()
	if err != nil {
		return nil, 0, fmt.Errorf("unable to open file: %v", err)
	}
	defer file.Close()

	writer := u.repo.NewWriter(repo.WriterOptions{
		Description: "FILE:" + f.Metadata().Name,
	})
	defer writer.Close()

	written, err := u.copyWithProgress(relativePath, writer, file, 0, f.Metadata().FileSize)
	if err != nil {
		u.Progress.Finished(relativePath, f.Metadata().FileSize, err)
		return nil, 0, err
	}

	e2, err := file.EntryMetadata()
	if err != nil {
		u.Progress.Finished(relativePath, f.Metadata().FileSize, err)
		return nil, 0, err
	}

	r, err := writer.Result(forceStored)
	if err != nil {
		u.Progress.Finished(relativePath, f.Metadata().FileSize, err)
		return nil, 0, err
	}

	de := newDirEntry(e2, r)
	de.FileSize = written

	u.Progress.Finished(relativePath, f.Metadata().FileSize, nil)

	return de, metadataHash(&de.EntryMetadata), nil
}

func (u *Uploader) uploaSymlinkInternal(f fs.Symlink, relativePath string) (*dir.Entry, uint64, error) {
	u.Progress.Started(relativePath, 1)

	target, err := f.Readlink()
	if err != nil {
		return nil, 0, fmt.Errorf("unable to read symlink: %v", err)
	}

	writer := u.repo.NewWriter(repo.WriterOptions{
		Description: "SYMLINK:" + f.Metadata().Name,
	})
	defer writer.Close()

	written, err := u.copyWithProgress(relativePath, writer, bytes.NewBufferString(target), 0, f.Metadata().FileSize)
	if err != nil {
		u.Progress.Finished(relativePath, f.Metadata().FileSize, err)
		return nil, 0, err
	}

	r, err := writer.Result(false)
	if err != nil {
		u.Progress.Finished(relativePath, f.Metadata().FileSize, err)
		return nil, 0, err
	}

	de := newDirEntry(f.Metadata(), r)
	de.FileSize = written
	u.Progress.Finished(relativePath, 1, nil)
	return de, metadataHash(&de.EntryMetadata), nil
}
func (u *Uploader) copyWithProgress(path string, dst io.Writer, src io.Reader, completed int64, length int64) (int64, error) {
	if u.uploadBuf == nil {
		u.uploadBuf = make([]byte, 128*1024) // 128 KB buffer
	}

	var written int64

	for {
		if u.IsCancelled() {
			return 0, errCancelled
		}

		readBytes, readErr := src.Read(u.uploadBuf)
		if readBytes > 0 {
			wroteBytes, writeErr := dst.Write(u.uploadBuf[0:readBytes])
			if wroteBytes > 0 {
				written += int64(wroteBytes)
				completed += int64(wroteBytes)
				if length < completed {
					length = completed
				}
				u.Progress.Progress(path, completed, length)
			}
			if writeErr != nil {
				return written, writeErr
			}
			if readBytes != wroteBytes {
				return written, io.ErrShortWrite
			}
		}

		if readErr != nil {
			if readErr == io.EOF {
				break
			}

			return written, readErr
		}
	}

	return written, nil
}

func newDirEntry(md *fs.EntryMetadata, oid repo.ObjectID) *dir.Entry {
	return &dir.Entry{
		EntryMetadata: *md,
		ObjectID:      oid,
	}
}

// uploadFile uploads the specified File to the repository.
func (u *Uploader) uploadFile(file fs.File) (repo.ObjectID, error) {
	e, _, err := u.uploadFileInternal(file, file.Metadata().Name, true)
	if err != nil {
		return repo.NullObjectID, err
	}
	return e.ObjectID, nil
}

// uploadDir uploads the specified Directory to the repository.
// An optional ID of a hash-cache object may be provided, in which case the Uploader will use its
// contents to avoid hashing
func (u *Uploader) uploadDir(dir fs.Directory) (repo.ObjectID, repo.ObjectID, error) {
	var err error

	if err := u.repo.BeginPacking(); err != nil {
		return repo.NullObjectID, repo.NullObjectID, err
	}

	mw := u.repo.NewWriter(repo.WriterOptions{
		Description:     "HASHCACHE:" + dir.Metadata().Name,
		BlockNamePrefix: "H",
		PackGroup:       "HC",
	})
	defer mw.Close()
	u.cacheWriter = hashcache.NewWriter(mw)
	oid, err := uploadDirInternal(u, dir, ".", true)
	if u.IsCancelled() {
		if err := u.cacheReader.CopyTo(u.cacheWriter); err != nil {
			return repo.NullObjectID, repo.NullObjectID, err
		}
	}
	u.cacheWriter.Finalize()
	u.cacheWriter = nil

	if err != nil {
		return repo.NullObjectID, repo.NullObjectID, err
	}

	hcid, err := mw.Result(true)
	if err := u.repo.FinishPacking(); err != nil {
		return repo.NullObjectID, repo.NullObjectID, fmt.Errorf("can't finish packing: %v", err)
	}
	return oid, hcid, err
}

func uploadDirInternal(
	u *Uploader,
	directory fs.Directory,
	relativePath string,
	forceStored bool,
) (repo.ObjectID, error) {
	u.Progress.StartedDir(relativePath)
	defer u.Progress.FinishedDir(relativePath)

	u.stats.TotalDirectoryCount++

	entries, err := directory.Readdir()
	if err != nil {
		return repo.NullObjectID, err
	}

	writer := u.repo.NewWriter(repo.WriterOptions{
		Description: "DIR:" + relativePath,
		PackGroup:   "DIR",
	})

	dw := dir.NewWriter(writer)
	defer writer.Close()

	for _, entry := range entries {
		if u.IsCancelled() {
			break
		}
		e := entry.Metadata()
		entryRelativePath := relativePath + "/" + e.Name

		var de *dir.Entry
		var hash uint64

		// regular file
		// See if we had this name during previous pass.
		cachedEntry := u.maybeIgnoreHashCacheEntry(u.cacheReader.FindEntry(entryRelativePath))

		// ... and whether file metadata is identical to the previous one.
		computedHash := metadataHash(e)
		cacheMatches := (cachedEntry != nil) && cachedEntry.Hash == computedHash

		switch entry.(type) {
		case fs.File:
			u.stats.TotalFileCount++
			u.stats.TotalFileSize += e.FileSize
		}

		if cacheMatches {
			u.stats.CachedFiles++
			u.Progress.Cached(entryRelativePath, entry.Metadata().FileSize)
			// Avoid hashing by reusing previous object ID.
			de, hash, err = newDirEntry(e, cachedEntry.ObjectID), cachedEntry.Hash, nil
		} else {
			switch entry := entry.(type) {
			case fs.Directory:
				var oid repo.ObjectID
				oid, err = uploadDirInternal(u, entry, entryRelativePath, false)
				de = newDirEntry(e, oid)
				hash = 0

			case fs.Symlink:
				de, hash, err = u.uploaSymlinkInternal(entry, entryRelativePath)

			case fs.File:
				u.stats.NonCachedFiles++
				de, hash, err = u.uploadFileInternal(entry, entryRelativePath, false)

			default:
				return repo.NullObjectID, fmt.Errorf("file type %v not supported", entry.Metadata().Type)
			}
		}

		if err == errCancelled {
			break
		}

		if err != nil {
			if u.IgnoreFileErrors {
				u.stats.ReadErrors++
				log.Printf("warning: unable to hash file %q: %s, ignoring", entryRelativePath, err)
				continue
			}
			return repo.NullObjectID, fmt.Errorf("unable to hash file: %s", err)
		}

		if err := dw.WriteEntry(de); err != nil {
			return repo.NullObjectID, err
		}

		if de.Type != fs.EntryTypeDirectory && hash != 0 {
			if err := u.cacheWriter.WriteEntry(hashcache.Entry{
				Name:     entryRelativePath,
				Hash:     hash,
				ObjectID: de.ObjectID,
			}); err != nil {
				return repo.NullObjectID, err
			}
		}
	}

	dw.Finalize()

	return writer.Result(forceStored)
}

func (u *Uploader) maybeIgnoreHashCacheEntry(e *hashcache.Entry) *hashcache.Entry {
	if rand.Intn(100) < u.ForceHashingPercentage {
		return nil
	}

	return e
}

// NewUploader creates new Uploader object for a given repository.
func NewUploader(r *repo.Repository) *Uploader {
	return &Uploader{
		repo:     r,
		Progress: &nullUploadProgress{},
	}
}

// Cancel requests cancellation of an upload that's in progress. Will typically result in an incomplete snapshot.
func (u *Uploader) Cancel() {
	atomic.StoreInt32(&u.cancelled, 1)
}

// Upload uploads contents of the specified filesystem entry (file or directory) to the repository and returns snapshot.Manifest with statistics.
// Old snapshot manifest, when provided can be used to speed up uploads by utilizing hash cache.
func (u *Uploader) Upload(
	source fs.Entry,
	sourceInfo *SourceInfo,
	old *Manifest,
) (*Manifest, error) {
	s := &Manifest{
		Source: *sourceInfo,
	}

	u.cacheReader = hashcache.Open(nil)
	u.stats = Stats{}
	if old != nil {
		if r, err := u.repo.Open(old.HashCacheID); err == nil {
			u.cacheReader = hashcache.Open(r)
		}
	}

	var err error

	s.StartTime = time.Now()

	switch entry := source.(type) {
	case fs.Directory:
		s.RootObjectID, s.HashCacheID, err = u.uploadDir(entry)

	case fs.File:
		s.RootObjectID, err = u.uploadFile(entry)

	default:
		return nil, fmt.Errorf("unsupported source: %v", s.Source)
	}
	if err != nil {
		return nil, err
	}

	s.IncompleteReason = u.cancelReason()
	s.EndTime = time.Now()
	s.Stats = u.stats
	s.Stats.Repository = u.repo.Status().Stats

	return s, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"reflect"
	"sort"
	"strings"

	"github.com/codegangsta/cli"
	mkr "github.com/mackerelio/mackerel-client-go"
	"github.com/mackerelio/mkr/logger"
)

var commandMonitors = cli.Command{
	Name:  "monitors",
	Usage: "Manipurate monitors",
	Description: `
    Manipurate monitor rules. Without subcommand, show all monitor rules.
    Request APIs under "/api/v0/monitors". See http://help-ja.mackerel.io/entry/spec/api/v0 .
`,
	Action: doMonitorsList,
	Subcommands: []cli.Command{
		{
			Name:        "pull",
			Usage:       "pull rules",
			Description: "Pull monitor rules from Mackerel server and save them to a file. The file can be specified by filepath argument <file>. The default is 'monitors.json'.",
			Action:      doMonitorsPull,
			Flags: []cli.Flag{
				cli.StringFlag{Name: "filepath, F", Value: "", Usage: "Filename to store monitor rule definitions. default: monitors.json"},
				cli.BoolFlag{Name: "verbose, v", Usage: "Verbose output mode"},
			},
		},
		{
			Name:        "diff",
			Usage:       "diff rules",
			Description: "Show difference of monitor rules between Mackerel and a file. The file can be specified by filepath argument <file>. The default is 'monitors.json'.",
			Action:      doMonitorsDiff,
			Flags: []cli.Flag{
				cli.StringFlag{Name: "filepath, F", Value: "", Usage: "Filename to store monitor rule definitions. default: monitors.json"},
			},
		},
		{
			Name:        "push",
			Usage:       "push rules",
			Description: "Push monitor rules, which are stored in a file, to Mackerel. The file can be specified by filepath argument <file>. The default is 'monitors.json'.",
			Action:      doMonitorsPush,
			Flags: []cli.Flag{
				cli.StringFlag{Name: "filepath, F", Value: "", Usage: "Filename to store monitor rule definitions. default: monitors.json"},
				cli.BoolFlag{Name: "dryrun, d", Usage: "Show which apis are called, but not execute."},
				cli.BoolFlag{Name: "verbose, v", Usage: "Verbose output mode"},
			},
		},
	},
}

func monitorSaveRules(rules []*(mkr.Monitor), optFilePath string) error {
	filePath := "monitors.json"
	if optFilePath != "" {
		filePath = optFilePath
	}
	file, err := os.Create(filePath)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	monitors := map[string]interface{}{"monitors": rules}
	dataRaw, err := json.MarshalIndent(monitors, "", "    ")
	logger.DieIf(err)
	data := strings.Replace(string(dataRaw), "\\u003c", "<", -1)
	data = strings.Replace(data, "\\u003e", ">", -1)

	_, err = file.WriteString(data)
	if err != nil {
		return err
	}
	return nil
}

func monitorLoadRules(optFilePath string) ([]*(mkr.Monitor), error) {
	filePath := "monitors.json"
	if optFilePath != "" {
		filePath = optFilePath
	}

	buff, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}

	var data struct {
		Monitors []*(mkr.Monitor) `json:"monitors"`
	}

	err = json.Unmarshal(buff, &data)
	if err != nil {
		return nil, err
	}
	return data.Monitors, nil
}

func doMonitorsList(c *cli.Context) {
	conffile := c.GlobalString("conf")

	monitors, err := newMackerel(conffile).FindMonitors()
	logger.DieIf(err)

	PrettyPrintJSON(monitors)
}

func doMonitorsPull(c *cli.Context) {
	conffile := c.GlobalString("conf")
	isVerbose := c.Bool("verbose")
	filePath := c.String("filepath")

	monitors, err := newMackerel(conffile).FindMonitors()
	logger.DieIf(err)

	monitorSaveRules(monitors, filePath)

	if isVerbose {
		PrettyPrintJSON(monitors)
	}
}

func appendDiff(src []string, name string, a interface{}, b interface{}) []string {
	diff := []string{}
	aType := reflect.TypeOf(a).String()
	format := "\"%s\""
	switch aType {
	case "uint64":
		format = "%d"
	case "float64":
		format = "%f"
	}
	if b != nil && a != b {
		diff = append(src, fmt.Sprintf("-  \"%s\": "+format+",", name, a))
		diff = append(diff, fmt.Sprintf("+  \"%s\": "+format+",", name, b))
	} else {
		diff = append(src, fmt.Sprintf("   \"%s\": "+format+",", name, a))
	}
	return diff
}

func printMonitor(a *mkr.Monitor, prefix string) {
	sA := reflect.ValueOf(a).Elem()
	diff := []string{" {"}
	for i := 0; i < sA.NumField(); i++ {
		fA := sA.Field(i)
		sAType := sA.Type()
		name := strings.Replace(sAType.Field(i).Tag.Get("json"), ",omitempty", "", 1)
		if sAType.Field(i).Type.String() != "[]string" {
			if name == "id" && fA.Interface() == "" {
				continue
			}
			diff = appendDiff(diff, name, fA.Interface(), nil)
		} else {
			diff = append(diff, fmt.Sprintf("   \"%s\": [", name))
			sortA := fA.Interface().([]string)
			sort.Strings(sortA)
			i := 0
			for i < len(sortA) {
				diff = append(diff, fmt.Sprintf("     \"%s\",", sortA[i]))
				i++
			}
			diff = append(diff, "   ],")
		}
	}
	diff = append(diff, " },")
	for _, d := range diff {
		fmt.Println(prefix + d)
	}
}

func diffMonitor(a *mkr.Monitor, b *mkr.Monitor) string {
	diff := []string{"  {"}
	diffNum := 0
	sA := reflect.ValueOf(a).Elem()
	sB := reflect.ValueOf(b).Elem()
	for i := 0; i < sA.NumField(); i++ {
		fA := sA.Field(i)
		fB := sB.Field(i)
		sAType := sA.Type()
		if sAType.Field(i).Type.String() != "[]string" {
			name := strings.Replace(sAType.Field(i).Tag.Get("json"), ",omitempty", "", 1)
			if name == "id" {
				continue
			}
			diff = appendDiff(diff, name, fA.Interface(), fB.Interface())
			if fA.Interface() != fB.Interface() {
				diffNum++
			}
		} else {
			name := strings.Replace(sAType.Field(i).Tag.Get("json"), ",omitempty", "", 1)
			diff = append(diff, fmt.Sprintf("    \"%s\": [", name))
			sortA := fA.Interface().([]string)
			sortB := fB.Interface().([]string)
			sort.Strings(sortA)
			sort.Strings(sortB)
			i := 0
			j := 0
			for i < len(sortA) && j < len(sortB) {
				if sortA[i] == sortB[j] {
					diff = append(diff, fmt.Sprintf("      \"%s\",", sortA[i]))
					i++
					j++
				} else if sortA[i] < sortB[j] {
					diff = append(diff, fmt.Sprintf("-     \"%s\",", sortA[i]))
					i++
					diffNum++
				} else if sortB[j] < sortA[i] {
					diff = append(diff, fmt.Sprintf("+     \"%s\",", sortB[j]))
					j++
					diffNum++
				}
			}
			diff = append(diff, "    ],")
		}
	}

	if diffNum > 0 {
		diff = append(diff, "  },")
		return strings.Join(diff, "\n")
	}
	return ""
}

func isSameMonitor(a *mkr.Monitor, b *mkr.Monitor, flagNameUniqueness bool) (string, bool) {
	if a == nil || b == nil {
		return "", false
	}
	if reflect.DeepEqual(*a, *b) {
		return "", true
	}
	if a.ID == b.ID || (flagNameUniqueness == true && b.ID == "" && a.Name == b.Name) {
		diff := diffMonitor(a, b)
		if diff != "" {
			return diff, false
		}
		return "", true
	}
	return "", false
}

func validateRules(monitors []*(mkr.Monitor), label string) (bool, error) {

	flagNameUniqueness := true
	// check each monitor
	for _, m := range monitors {
		v := reflect.ValueOf(m).Elem()
		for _, f := range []string{"Type"} {
			vf := v.FieldByName(f)
			if !vf.IsValid() || (vf.Type().String() == "string" && vf.Interface() == "") {
				return false, fmt.Errorf("Monitor '%s' should have '%s': %s", label, f, v.FieldByName(f).Interface())
			}
		}
		switch m.Type {
		case "host", "service":
			for _, f := range []string{"Name", "Metric"} {
				vf := v.FieldByName(f)
				if !vf.IsValid() || (vf.Type().String() == "string" && vf.Interface() == "") {
					return false, fmt.Errorf("Monitor '%s' should have '%s': %s", label, f, v.FieldByName(f).Interface())
				}
			}
		case "external":
			for _, f := range []string{"Name", "URL"} {
				vf := v.FieldByName(f)
				if !vf.IsValid() || (vf.Type().String() == "string" && vf.Interface() == "") {
					return false, fmt.Errorf("Monitor '%s' should have '%s': %s", label, f, v.FieldByName(f).Interface())
				}
			}
		case "connectivity":
		default:
			return false, fmt.Errorf("Unknown type is found: %s", m.Type)
		}
	}

	// check name uniqueness
	names := map[string]bool{}
	for _, m := range monitors {
		if names[m.Name] {
			logger.Log("Warning: ", fmt.Sprintf("Names of %s is not unique.", label))
			flagNameUniqueness = false
		}
		names[m.Name] = true
	}
	return flagNameUniqueness, nil
}

type monitorDiffPair struct {
	remote *mkr.Monitor
	local  *mkr.Monitor
}

type monitorDiff struct {
	onlyRemote []*(mkr.Monitor)
	onlyLocal  []*(mkr.Monitor)
	diff       []*monitorDiffPair
}

func checkMonitorsDiff(c *cli.Context) monitorDiff {
	conffile := c.GlobalString("conf")
	filePath := c.String("filepath")

	var monitorDiff monitorDiff

	monitorsRemote, err := newMackerel(conffile).FindMonitors()
	logger.DieIf(err)
	flagNameUniquenessRemote, err := validateRules(monitorsRemote, "remote rules")
	logger.DieIf(err)

	monitorsLocal, err := monitorLoadRules(filePath)
	logger.DieIf(err)
	flagNameUniquenessLocal, err := validateRules(monitorsLocal, "local rules")
	logger.DieIf(err)

	flagNameUniqueness := flagNameUniquenessLocal && flagNameUniquenessRemote

	for _, remote := range monitorsRemote {
		found := false
		for i, local := range monitorsLocal {
			diff, isSame := isSameMonitor(remote, local, flagNameUniqueness)
			if isSame || diff != "" {
				monitorsLocal[i] = nil
				found = true
				if diff != "" {
					monitorDiff.diff = append(monitorDiff.diff, &monitorDiffPair{remote, local})
				}
				break
			}
		}
		if found == false {
			monitorDiff.onlyRemote = append(monitorDiff.onlyRemote, remote)
		}
	}
	for _, local := range monitorsLocal {
		if local != nil {
			monitorDiff.onlyLocal = append(monitorDiff.onlyLocal, local)
		}
	}

	return monitorDiff
}

func doMonitorsDiff(c *cli.Context) {
	monitorDiff := checkMonitorsDiff(c)

	var diffs []string
	for _, d := range monitorDiff.diff {
		diffs = append(diffs, diffMonitor(d.remote, d.local))
	}

	fmt.Printf("Summary: %d modify, %d append, %d remove\n\n", len(monitorDiff.diff), len(monitorDiff.onlyLocal), len(monitorDiff.onlyRemote))
	for _, diff := range diffs {
		fmt.Println(diff)
	}
	for _, m := range monitorDiff.onlyRemote {
		printMonitor(m, "-")
	}
	for _, m := range monitorDiff.onlyLocal {
		printMonitor(m, "+")
	}
}

func doMonitorsPush(c *cli.Context) {
	monitorDiff := checkMonitorsDiff(c)
	isDryRun := c.Bool("dryRun")
	isVerbose := c.Bool("verbose")

	conffile := c.GlobalString("conf")
	client := newMackerel(conffile)
	if isVerbose {
		client.Verbose = true
	}

	for _, m := range monitorDiff.onlyLocal {
		logger.Log("info", "Create a new rule.")
		printMonitor(m, "")
		if !isDryRun {
			_, err := client.CreateMonitor(m)
			logger.DieIf(err)
		}
	}
	for _, m := range monitorDiff.onlyRemote {
		logger.Log("info", "Delete a rule.")
		printMonitor(m, "")
		if !isDryRun {
			_, err := client.DeleteMonitor(m.ID)
			logger.DieIf(err)
		}
	}
	for _, d := range monitorDiff.diff {
		logger.Log("info", "Update a rule.")
		printMonitor(d.local, "")
		if !isDryRun {
			_, err := client.UpdateMonitor(d.remote.ID, d.local)
			logger.DieIf(err)
		}
	}

}
<commit_msg>dryrun now works well<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"reflect"
	"sort"
	"strings"

	"github.com/codegangsta/cli"
	mkr "github.com/mackerelio/mackerel-client-go"
	"github.com/mackerelio/mkr/logger"
)

var commandMonitors = cli.Command{
	Name:  "monitors",
	Usage: "Manipurate monitors",
	Description: `
    Manipurate monitor rules. Without subcommand, show all monitor rules.
    Request APIs under "/api/v0/monitors". See http://help-ja.mackerel.io/entry/spec/api/v0 .
`,
	Action: doMonitorsList,
	Subcommands: []cli.Command{
		{
			Name:        "pull",
			Usage:       "pull rules",
			Description: "Pull monitor rules from Mackerel server and save them to a file. The file can be specified by filepath argument <file>. The default is 'monitors.json'.",
			Action:      doMonitorsPull,
			Flags: []cli.Flag{
				cli.StringFlag{Name: "filepath, F", Value: "", Usage: "Filename to store monitor rule definitions. default: monitors.json"},
				cli.BoolFlag{Name: "verbose, v", Usage: "Verbose output mode"},
			},
		},
		{
			Name:        "diff",
			Usage:       "diff rules",
			Description: "Show difference of monitor rules between Mackerel and a file. The file can be specified by filepath argument <file>. The default is 'monitors.json'.",
			Action:      doMonitorsDiff,
			Flags: []cli.Flag{
				cli.StringFlag{Name: "filepath, F", Value: "", Usage: "Filename to store monitor rule definitions. default: monitors.json"},
			},
		},
		{
			Name:        "push",
			Usage:       "push rules",
			Description: "Push monitor rules, which are stored in a file, to Mackerel. The file can be specified by filepath argument <file>. The default is 'monitors.json'.",
			Action:      doMonitorsPush,
			Flags: []cli.Flag{
				cli.StringFlag{Name: "filepath, F", Value: "", Usage: "Filename to store monitor rule definitions. default: monitors.json"},
				cli.BoolFlag{Name: "dryrun, d", Usage: "Show which apis are called, but not execute."},
				cli.BoolFlag{Name: "verbose, v", Usage: "Verbose output mode"},
			},
		},
	},
}

func monitorSaveRules(rules []*(mkr.Monitor), optFilePath string) error {
	filePath := "monitors.json"
	if optFilePath != "" {
		filePath = optFilePath
	}
	file, err := os.Create(filePath)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	monitors := map[string]interface{}{"monitors": rules}
	dataRaw, err := json.MarshalIndent(monitors, "", "    ")
	logger.DieIf(err)
	data := strings.Replace(string(dataRaw), "\\u003c", "<", -1)
	data = strings.Replace(data, "\\u003e", ">", -1)

	_, err = file.WriteString(data)
	if err != nil {
		return err
	}
	return nil
}

func monitorLoadRules(optFilePath string) ([]*(mkr.Monitor), error) {
	filePath := "monitors.json"
	if optFilePath != "" {
		filePath = optFilePath
	}

	buff, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}

	var data struct {
		Monitors []*(mkr.Monitor) `json:"monitors"`
	}

	err = json.Unmarshal(buff, &data)
	if err != nil {
		return nil, err
	}
	return data.Monitors, nil
}

func doMonitorsList(c *cli.Context) {
	conffile := c.GlobalString("conf")

	monitors, err := newMackerel(conffile).FindMonitors()
	logger.DieIf(err)

	PrettyPrintJSON(monitors)
}

func doMonitorsPull(c *cli.Context) {
	conffile := c.GlobalString("conf")
	isVerbose := c.Bool("verbose")
	filePath := c.String("filepath")

	monitors, err := newMackerel(conffile).FindMonitors()
	logger.DieIf(err)

	monitorSaveRules(monitors, filePath)

	if isVerbose {
		PrettyPrintJSON(monitors)
	}
}

func appendDiff(src []string, name string, a interface{}, b interface{}) []string {
	diff := []string{}
	aType := reflect.TypeOf(a).String()
	format := "\"%s\""
	switch aType {
	case "uint64":
		format = "%d"
	case "float64":
		format = "%f"
	}
	if b != nil && a != b {
		diff = append(src, fmt.Sprintf("-  \"%s\": "+format+",", name, a))
		diff = append(diff, fmt.Sprintf("+  \"%s\": "+format+",", name, b))
	} else {
		diff = append(src, fmt.Sprintf("   \"%s\": "+format+",", name, a))
	}
	return diff
}

func printMonitor(a *mkr.Monitor, prefix string) {
	sA := reflect.ValueOf(a).Elem()
	diff := []string{" {"}
	for i := 0; i < sA.NumField(); i++ {
		fA := sA.Field(i)
		sAType := sA.Type()
		name := strings.Replace(sAType.Field(i).Tag.Get("json"), ",omitempty", "", 1)
		if sAType.Field(i).Type.String() != "[]string" {
			if name == "id" && fA.Interface() == "" {
				continue
			}
			diff = appendDiff(diff, name, fA.Interface(), nil)
		} else {
			diff = append(diff, fmt.Sprintf("   \"%s\": [", name))
			sortA := fA.Interface().([]string)
			sort.Strings(sortA)
			i := 0
			for i < len(sortA) {
				diff = append(diff, fmt.Sprintf("     \"%s\",", sortA[i]))
				i++
			}
			diff = append(diff, "   ],")
		}
	}
	diff = append(diff, " },")
	for _, d := range diff {
		fmt.Println(prefix + d)
	}
}

func diffMonitor(a *mkr.Monitor, b *mkr.Monitor) string {
	diff := []string{"  {"}
	diffNum := 0
	sA := reflect.ValueOf(a).Elem()
	sB := reflect.ValueOf(b).Elem()
	for i := 0; i < sA.NumField(); i++ {
		fA := sA.Field(i)
		fB := sB.Field(i)
		sAType := sA.Type()
		if sAType.Field(i).Type.String() != "[]string" {
			name := strings.Replace(sAType.Field(i).Tag.Get("json"), ",omitempty", "", 1)
			if name == "id" {
				continue
			}
			diff = appendDiff(diff, name, fA.Interface(), fB.Interface())
			if fA.Interface() != fB.Interface() {
				diffNum++
			}
		} else {
			name := strings.Replace(sAType.Field(i).Tag.Get("json"), ",omitempty", "", 1)
			diff = append(diff, fmt.Sprintf("    \"%s\": [", name))
			sortA := fA.Interface().([]string)
			sortB := fB.Interface().([]string)
			sort.Strings(sortA)
			sort.Strings(sortB)
			i := 0
			j := 0
			for i < len(sortA) && j < len(sortB) {
				if sortA[i] == sortB[j] {
					diff = append(diff, fmt.Sprintf("      \"%s\",", sortA[i]))
					i++
					j++
				} else if sortA[i] < sortB[j] {
					diff = append(diff, fmt.Sprintf("-     \"%s\",", sortA[i]))
					i++
					diffNum++
				} else if sortB[j] < sortA[i] {
					diff = append(diff, fmt.Sprintf("+     \"%s\",", sortB[j]))
					j++
					diffNum++
				}
			}
			diff = append(diff, "    ],")
		}
	}

	if diffNum > 0 {
		diff = append(diff, "  },")
		return strings.Join(diff, "\n")
	}
	return ""
}

func isSameMonitor(a *mkr.Monitor, b *mkr.Monitor, flagNameUniqueness bool) (string, bool) {
	if a == nil || b == nil {
		return "", false
	}
	if reflect.DeepEqual(*a, *b) {
		return "", true
	}
	if a.ID == b.ID || (flagNameUniqueness == true && b.ID == "" && a.Name == b.Name) {
		diff := diffMonitor(a, b)
		if diff != "" {
			return diff, false
		}
		return "", true
	}
	return "", false
}

func validateRules(monitors []*(mkr.Monitor), label string) (bool, error) {

	flagNameUniqueness := true
	// check each monitor
	for _, m := range monitors {
		v := reflect.ValueOf(m).Elem()
		for _, f := range []string{"Type"} {
			vf := v.FieldByName(f)
			if !vf.IsValid() || (vf.Type().String() == "string" && vf.Interface() == "") {
				return false, fmt.Errorf("Monitor '%s' should have '%s': %s", label, f, v.FieldByName(f).Interface())
			}
		}
		switch m.Type {
		case "host", "service":
			for _, f := range []string{"Name", "Metric"} {
				vf := v.FieldByName(f)
				if !vf.IsValid() || (vf.Type().String() == "string" && vf.Interface() == "") {
					return false, fmt.Errorf("Monitor '%s' should have '%s': %s", label, f, v.FieldByName(f).Interface())
				}
			}
		case "external":
			for _, f := range []string{"Name", "URL"} {
				vf := v.FieldByName(f)
				if !vf.IsValid() || (vf.Type().String() == "string" && vf.Interface() == "") {
					return false, fmt.Errorf("Monitor '%s' should have '%s': %s", label, f, v.FieldByName(f).Interface())
				}
			}
		case "connectivity":
		default:
			return false, fmt.Errorf("Unknown type is found: %s", m.Type)
		}
	}

	// check name uniqueness
	names := map[string]bool{}
	for _, m := range monitors {
		if names[m.Name] {
			logger.Log("Warning: ", fmt.Sprintf("Names of %s is not unique.", label))
			flagNameUniqueness = false
		}
		names[m.Name] = true
	}
	return flagNameUniqueness, nil
}

type monitorDiffPair struct {
	remote *mkr.Monitor
	local  *mkr.Monitor
}

type monitorDiff struct {
	onlyRemote []*(mkr.Monitor)
	onlyLocal  []*(mkr.Monitor)
	diff       []*monitorDiffPair
}

func checkMonitorsDiff(c *cli.Context) monitorDiff {
	conffile := c.GlobalString("conf")
	filePath := c.String("filepath")

	var monitorDiff monitorDiff

	monitorsRemote, err := newMackerel(conffile).FindMonitors()
	logger.DieIf(err)
	flagNameUniquenessRemote, err := validateRules(monitorsRemote, "remote rules")
	logger.DieIf(err)

	monitorsLocal, err := monitorLoadRules(filePath)
	logger.DieIf(err)
	flagNameUniquenessLocal, err := validateRules(monitorsLocal, "local rules")
	logger.DieIf(err)

	flagNameUniqueness := flagNameUniquenessLocal && flagNameUniquenessRemote

	for _, remote := range monitorsRemote {
		found := false
		for i, local := range monitorsLocal {
			diff, isSame := isSameMonitor(remote, local, flagNameUniqueness)
			if isSame || diff != "" {
				monitorsLocal[i] = nil
				found = true
				if diff != "" {
					monitorDiff.diff = append(monitorDiff.diff, &monitorDiffPair{remote, local})
				}
				break
			}
		}
		if found == false {
			monitorDiff.onlyRemote = append(monitorDiff.onlyRemote, remote)
		}
	}
	for _, local := range monitorsLocal {
		if local != nil {
			monitorDiff.onlyLocal = append(monitorDiff.onlyLocal, local)
		}
	}

	return monitorDiff
}

func doMonitorsDiff(c *cli.Context) {
	monitorDiff := checkMonitorsDiff(c)

	var diffs []string
	for _, d := range monitorDiff.diff {
		diffs = append(diffs, diffMonitor(d.remote, d.local))
	}

	fmt.Printf("Summary: %d modify, %d append, %d remove\n\n", len(monitorDiff.diff), len(monitorDiff.onlyLocal), len(monitorDiff.onlyRemote))
	for _, diff := range diffs {
		fmt.Println(diff)
	}
	for _, m := range monitorDiff.onlyRemote {
		printMonitor(m, "-")
	}
	for _, m := range monitorDiff.onlyLocal {
		printMonitor(m, "+")
	}
}

func doMonitorsPush(c *cli.Context) {
	monitorDiff := checkMonitorsDiff(c)
	isDryRun := c.Bool("dryrun")
	isVerbose := c.Bool("verbose")

	conffile := c.GlobalString("conf")
	client := newMackerel(conffile)
	if isVerbose {
		client.Verbose = true
	}

	for _, m := range monitorDiff.onlyLocal {
		logger.Log("info", "Create a new rule.")
		printMonitor(m, "")
		if !isDryRun {
			_, err := client.CreateMonitor(m)
			logger.DieIf(err)
		}
	}
	for _, m := range monitorDiff.onlyRemote {
		logger.Log("info", "Delete a rule.")
		printMonitor(m, "")
		if !isDryRun {
			_, err := client.DeleteMonitor(m.ID)
			logger.DieIf(err)
		}
	}
	for _, d := range monitorDiff.diff {
		logger.Log("info", "Update a rule.")
		printMonitor(d.local, "")
		if !isDryRun {
			_, err := client.UpdateMonitor(d.remote.ID, d.local)
			logger.DieIf(err)
		}
	}

}
<|endoftext|>


<commit_before>package mqtt

// Values of QoS.
const (
	QoS0 = iota
	QoS1
	QoS2
)
<commit_msg>Update mqtt/qos.go<commit_after>package mqtt

// Values of QoS.
const (
	// QoS0 represents "QoS 0: At most once delivery".
	QoS0 = iota
	// QoS1 represents "QoS 1: At least once delivery".
	QoS1
	// QoS2 represents "QoS 2: Exactly once delivery".
	QoS2
)
<|endoftext|>


<commit_before>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package msg

import (
	"bytes"
	"crypto/sha512"
	"io"

	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/msg/session"
	"github.com/mutecomm/mute/util/bzero"
	"golang.org/x/crypto/hkdf"
)

// checkKeys checks that the keys k1, k2, k3, and k4 are pairwise different to
// prevent possible reflection attacks.
func checkKeys(k1, k2, k3, k4 *[32]byte) error {
	if bytes.Equal(k1[:], k2[:]) {
		return ErrReflection
	}
	if bytes.Equal(k1[:], k3[:]) {
		return ErrReflection
	}
	if bytes.Equal(k1[:], k4[:]) {
		return ErrReflection
	}
	if bytes.Equal(k2[:], k3[:]) {
		return ErrReflection
	}
	if bytes.Equal(k2[:], k4[:]) {
		return ErrReflection
	}
	if bytes.Equal(k3[:], k4[:]) {
		return ErrReflection
	}
	return nil
}

// deriveRootKey derives the next root key from t1, t2, t3, and the
// previousRootKeyHash (if it exists).
func deriveRootKey(
	t1, t2, t3 *[32]byte,
	previousRootKeyHash *[64]byte,
) (*[24]byte, error) {
	master := make([]byte, 32+32+32+64)
	copy(master[:], t1[:])
	copy(master[32:], t2[:])
	copy(master[64:], t3[:])
	if previousRootKeyHash != nil {
		copy(master[96:], previousRootKeyHash[:])
	}

	hkdf := hkdf.New(sha512.New, master, nil, nil)

	// derive root key
	// TODO: size correct? Shouldn't it be 64 bytes?
	var rootKey [24]byte
	if _, err := io.ReadFull(hkdf, rootKey[:]); err != nil {
		return nil, err
	}

	return &rootKey, nil
}

// generateMessageKeys generates the next numOfKeys many session keys from
// from rootKey for given senderIdentity and recipientIdentity.
// If recipientKeys is true the generated sender and reciever keys are stored in
// reverse order.
// It uses senderSessionPub and recipientPub in the process and calls
// keyStore.StoresSession and keyStore.SetSessionState to store the result.
func generateMessageKeys(
	senderIdentity, recipientIdentity string,
	rootKey *[24]byte,
	recipientKeys bool,
	senderSessionPub, recipientPub *[32]byte,
	numOfKeys int,
	keyStore session.Store,
) error {
	var (
		identities string
		send       []string
		recv       []string
	)

	// identity_fix = HASH(SORT(SenderNym, RecipientNym))
	if senderIdentity < recipientIdentity {
		identities = senderIdentity + recipientIdentity
	} else {
		identities = recipientIdentity + senderIdentity
	}
	identityFix := cipher.SHA512([]byte(identities))

	chainKey := rootKey[:]
	for i := 0; i < numOfKeys; i++ {
		// messagekey_send[i] = HMAC_HASH(chainkey, "MESSAGE" | HASH(RecipientPub) | identity_fix)
		buffer := append([]byte("MESSAGE"), cipher.SHA512(recipientPub[:])...)
		buffer = append(buffer, identityFix...)
		send = append(send, base64.Encode(cipher.HMAC(chainKey, buffer)))

		// messagekey_recv[i] = HMAC_HASH(chainkey, "MESSAGE" | HASH(SenderSessionPub) | identity_fix)
		buffer = append([]byte("MESSAGE"), cipher.SHA512(senderSessionPub[:])...)
		buffer = append(buffer, identityFix...)
		recv = append(recv, base64.Encode(cipher.HMAC(chainKey, buffer)))

		// chainkey = HMAC_HASH(chainkey, "CHAIN" )
		chainKey = cipher.HMAC(chainKey, []byte("CHAIN"))
	}

	// calculate root key hash
	rootKeyHash := base64.Encode(cipher.SHA512(rootKey[:]))
	bzero.Bytes(rootKey[:])

	// reverse key material, if necessary
	if recipientKeys {
		senderIdentity, recipientIdentity = recipientIdentity, senderIdentity
		send, recv = recv, send
	}

	// store session
	var senderSessionPubHash string
	if recipientKeys {
		senderSessionPubHash = base64.Encode(cipher.SHA512(recipientPub[:]))
	} else {
		senderSessionPubHash = base64.Encode(cipher.SHA512(senderSessionPub[:]))
	}
	err := keyStore.StoreSession(senderIdentity, recipientIdentity,
		senderSessionPubHash, rootKeyHash, base64.Encode(chainKey), send, recv)
	if err != nil {
		return err
	}

	return nil
}

// deriveSymmetricKeys derives the symmetric cryptoKey and hmacKey from the
// given messageKey.
func deriveSymmetricKeys(messageKey *[64]byte) (
	cryptoKey, hmacKey []byte,
	err error,
) {
	// TODO: set optional salt and info?
	hkdf := hkdf.New(sha512.New, messageKey[:], nil, nil)

	// derive crypto key for AES-256
	cryptoKey = make([]byte, 32)
	if _, err := io.ReadFull(hkdf, cryptoKey); err != nil {
		return nil, nil, err
	}

	// derive HMAC key for SHA-512 HMAC (TODO: correct size?)
	hmacKey = make([]byte, 64)
	if _, err := io.ReadFull(hkdf, hmacKey); err != nil {
		return nil, nil, err
	}

	return
}
<commit_msg>msg: make generateMessageKeys() more efficient<commit_after>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package msg

import (
	"bytes"
	"crypto/sha512"
	"io"

	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/msg/session"
	"github.com/mutecomm/mute/util/bzero"
	"golang.org/x/crypto/hkdf"
)

// checkKeys checks that the keys k1, k2, k3, and k4 are pairwise different to
// prevent possible reflection attacks.
func checkKeys(k1, k2, k3, k4 *[32]byte) error {
	if bytes.Equal(k1[:], k2[:]) {
		return ErrReflection
	}
	if bytes.Equal(k1[:], k3[:]) {
		return ErrReflection
	}
	if bytes.Equal(k1[:], k4[:]) {
		return ErrReflection
	}
	if bytes.Equal(k2[:], k3[:]) {
		return ErrReflection
	}
	if bytes.Equal(k2[:], k4[:]) {
		return ErrReflection
	}
	if bytes.Equal(k3[:], k4[:]) {
		return ErrReflection
	}
	return nil
}

// deriveRootKey derives the next root key from t1, t2, t3, and the
// previousRootKeyHash (if it exists).
func deriveRootKey(
	t1, t2, t3 *[32]byte,
	previousRootKeyHash *[64]byte,
) (*[24]byte, error) {
	master := make([]byte, 32+32+32+64)
	copy(master[:], t1[:])
	copy(master[32:], t2[:])
	copy(master[64:], t3[:])
	if previousRootKeyHash != nil {
		copy(master[96:], previousRootKeyHash[:])
	}

	hkdf := hkdf.New(sha512.New, master, nil, nil)

	// derive root key
	// TODO: size correct? Shouldn't it be 64 bytes?
	var rootKey [24]byte
	if _, err := io.ReadFull(hkdf, rootKey[:]); err != nil {
		return nil, err
	}

	return &rootKey, nil
}

// generateMessageKeys generates the next numOfKeys many session keys from
// from rootKey for given senderIdentity and recipientIdentity.
// If recipientKeys is true the generated sender and reciever keys are stored in
// reverse order.
// It uses senderSessionPub and recipientPub in the process and calls
// keyStore.StoresSession and keyStore.SetSessionState to store the result.
func generateMessageKeys(
	senderIdentity, recipientIdentity string,
	rootKey *[24]byte,
	recipientKeys bool,
	senderSessionPub, recipientPub *[32]byte,
	numOfKeys int,
	keyStore session.Store,
) error {
	var (
		identities string
		send       []string
		recv       []string
	)

	// identity_fix = HASH(SORT(SenderNym, RecipientNym))
	if senderIdentity < recipientIdentity {
		identities = senderIdentity + recipientIdentity
	} else {
		identities = recipientIdentity + senderIdentity
	}
	identityFix := cipher.SHA512([]byte(identities))
	recipientPubHash := cipher.SHA512(recipientPub[:])
	senderSessionPubHash := cipher.SHA512(senderSessionPub[:])

	chainKey := rootKey[:]
	for i := 0; i < numOfKeys; i++ {
		// messagekey_send[i] = HMAC_HASH(chainkey, "MESSAGE" | HASH(RecipientPub) | identity_fix)
		buffer := append([]byte("MESSAGE"), recipientPubHash...)
		buffer = append(buffer, identityFix...)
		send = append(send, base64.Encode(cipher.HMAC(chainKey, buffer)))

		// messagekey_recv[i] = HMAC_HASH(chainkey, "MESSAGE" | HASH(SenderSessionPub) | identity_fix)
		buffer = append([]byte("MESSAGE"), senderSessionPubHash...)
		buffer = append(buffer, identityFix...)
		recv = append(recv, base64.Encode(cipher.HMAC(chainKey, buffer)))

		// chainkey = HMAC_HASH(chainkey, "CHAIN" )
		chainKey = cipher.HMAC(chainKey, []byte("CHAIN"))
	}

	// calculate root key hash
	rootKeyHash := base64.Encode(cipher.SHA512(rootKey[:]))
	bzero.Bytes(rootKey[:])

	// reverse key material, if necessary
	if recipientKeys {
		senderIdentity, recipientIdentity = recipientIdentity, senderIdentity
		send, recv = recv, send
	}

	// store session
	var pubHash string
	if recipientKeys {
		pubHash = base64.Encode(cipher.SHA512(recipientPub[:]))
	} else {
		pubHash = base64.Encode(cipher.SHA512(senderSessionPub[:]))
	}
	err := keyStore.StoreSession(senderIdentity, recipientIdentity, pubHash,
		rootKeyHash, base64.Encode(chainKey), send, recv)
	if err != nil {
		return err
	}

	return nil
}

// deriveSymmetricKeys derives the symmetric cryptoKey and hmacKey from the
// given messageKey.
func deriveSymmetricKeys(messageKey *[64]byte) (
	cryptoKey, hmacKey []byte,
	err error,
) {
	// TODO: set optional salt and info?
	hkdf := hkdf.New(sha512.New, messageKey[:], nil, nil)

	// derive crypto key for AES-256
	cryptoKey = make([]byte, 32)
	if _, err := io.ReadFull(hkdf, cryptoKey); err != nil {
		return nil, nil, err
	}

	// derive HMAC key for SHA-512 HMAC (TODO: correct size?)
	hmacKey = make([]byte, 64)
	if _, err := io.ReadFull(hkdf, hmacKey); err != nil {
		return nil, nil, err
	}

	return
}
<|endoftext|>


<commit_before>//usr/bin/env go run $0 $@; exit;
package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/cbroglie/mustache"
	"github.com/docopt/docopt-go"
	"github.com/ghodss/yaml"
)

var logger = log.New(os.Stderr, "", 0)

func main() {
	doc := `Mustache Cli

        Command line interface for rendering mustache templates.
        Data is either expected via data option with a file name or
		via stdin. If data option is given that will be used.

        Examples:
            # Basic template usage
            mustache data.json template.mustache

            # Pull variables from environment
            mustache ENV template.mustache

            # Pull variables from environment with overrides. This will merge starting with env vars.
            # Think of order as priority.
			mustache ENV template.mustache --override data.json --override data1.json

			# get base data from stdin
    		cat data-source.json | mustache template.mustache

        Usage:
            mustache [<data-file>] <template-path> [--override=<data-file>]...
            mustache <template-path> [--override=<data-file>]...

        Arguments:
            <data-file>      Path to data file. ENV is a special identifier to use environment variables.

			<template-path>  Path to template file.

        Options:
            -h --help            Show help message.

            -o --override <file> Override data files. Overrides will be done in order.
	`

	arguments, _ := docopt.Parse(doc, nil, true, "Mustache 1.0.0", false)
	dataPath := arguments["<data-file>"]
	templatePath := arguments["<template-path>"].(string)
	overrideList := arguments["--override"].([]string)

	var (
		err     error
		data    interface{}
		context = make([]interface{}, 1+len(overrideList))
	)

	if dataPath == nil {
		data, err = loadFromStdin()
	} else {
		path := dataPath.(string)
		data = loadFromEnvOrFile(path)
	}
	context[0] = data

	for i, override := range overrideList {
		newData := loadFromEnvOrFile(override)
		checkErr(err)
		context[i+1] = newData
	}

	checkErr(err)
	output, err := mustache.RenderFile(templatePath, context...)
	checkErr(err)
	fmt.Println(output)
}

func checkErr(err error) {
	if err != nil {
		logError("Error occurred rendering template", err)
		os.Exit(1)
	}
}

func loadFromEnvOrFile(path string) interface{} {
	if path == "ENV" {
		return loadFromEnv()
	}

	d, err := loadFromFile(path)
	checkErr(err)
	return d
}

func loadFromEnv() interface{} {
	m := map[string]string{}

	for _, item := range os.Environ() {
		splits := strings.Split(item, "=")
		m[splits[0]] = os.Getenv(splits[0])
	}

	return m
}

func loadFromFile(path string) (interface{}, error) {
	raw, readErr := ioutil.ReadFile(path)

	if readErr != nil {
		return nil, readErr
	}

	return decodeData(raw)
}

func loadFromStdin() (interface{}, error) {
	raw, readErr := ioutil.ReadAll(os.Stdin)

	if readErr != nil {
		return nil, readErr
	}

	return decodeData(raw)
}

func decodeData(raw []byte) (interface{}, error) {
	var data interface{}
	err := yaml.Unmarshal(raw, &data)

	return data, err
}

func logError(msg string, err error) {
	logger.Println(msg)
	logger.Println(err.Error())
}
<commit_msg>Don't print trailing newline<commit_after>//usr/bin/env go run $0 $@; exit;
package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/cbroglie/mustache"
	"github.com/docopt/docopt-go"
	"github.com/ghodss/yaml"
)

var logger = log.New(os.Stderr, "", 0)

func main() {
	doc := `Mustache Cli

        Command line interface for rendering mustache templates.
        Data is either expected via data option with a file name or
		via stdin. If data option is given that will be used.

        Examples:
            # Basic template usage
            mustache data.json template.mustache

            # Pull variables from environment
            mustache ENV template.mustache

            # Pull variables from environment with overrides. This will merge starting with env vars.
            # Think of order as priority.
			mustache ENV template.mustache --override data.json --override data1.json

			# get base data from stdin
    		cat data-source.json | mustache template.mustache

        Usage:
            mustache [<data-file>] <template-path> [--override=<data-file>]...
            mustache <template-path> [--override=<data-file>]...

        Arguments:
            <data-file>      Path to data file. ENV is a special identifier to use environment variables.

			<template-path>  Path to template file.

        Options:
            -h --help            Show help message.

            -o --override <file> Override data files. Overrides will be done in order.
	`

	arguments, _ := docopt.Parse(doc, nil, true, "Mustache 1.0.0", false)
	dataPath := arguments["<data-file>"]
	templatePath := arguments["<template-path>"].(string)
	overrideList := arguments["--override"].([]string)

	var (
		err     error
		data    interface{}
		context = make([]interface{}, 1+len(overrideList))
	)

	if dataPath == nil {
		data, err = loadFromStdin()
	} else {
		path := dataPath.(string)
		data = loadFromEnvOrFile(path)
	}
	context[0] = data

	for i, override := range overrideList {
		newData := loadFromEnvOrFile(override)
		checkErr(err)
		context[i+1] = newData
	}

	checkErr(err)
	output, err := mustache.RenderFile(templatePath, context...)
	checkErr(err)
	fmt.Print(output)
}

func checkErr(err error) {
	if err != nil {
		logError("Error occurred rendering template", err)
		os.Exit(1)
	}
}

func loadFromEnvOrFile(path string) interface{} {
	if path == "ENV" {
		return loadFromEnv()
	}

	d, err := loadFromFile(path)
	checkErr(err)
	return d
}

func loadFromEnv() interface{} {
	m := map[string]string{}

	for _, item := range os.Environ() {
		splits := strings.Split(item, "=")
		m[splits[0]] = os.Getenv(splits[0])
	}

	return m
}

func loadFromFile(path string) (interface{}, error) {
	raw, readErr := ioutil.ReadFile(path)

	if readErr != nil {
		return nil, readErr
	}

	return decodeData(raw)
}

func loadFromStdin() (interface{}, error) {
	raw, readErr := ioutil.ReadAll(os.Stdin)

	if readErr != nil {
		return nil, readErr
	}

	return decodeData(raw)
}

func decodeData(raw []byte) (interface{}, error) {
	var data interface{}
	err := yaml.Unmarshal(raw, &data)

	return data, err
}

func logError(msg string, err error) {
	logger.Println(msg)
	logger.Println(err.Error())
}
<|endoftext|>


<commit_before>package mstate

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"launchpad.net/juju-core/charm"
	"sort"
	"strings"
)

// RelationRole defines the role of a relation endpoint.
type RelationRole string

const (
	RoleProvider RelationRole = "provider"
	RoleRequirer RelationRole = "requirer"
	RolePeer     RelationRole = "peer"
)

// counterpartRole returns the RelationRole that this RelationRole
// can relate to.
// This should remain an internal method because the relation
// model does not guarantee that for every role there will
// necessarily exist a single counterpart role that is sensible
// for basing algorithms upon.
func (r RelationRole) counterpartRole() RelationRole {
	switch r {
	case RoleProvider:
		return RoleRequirer
	case RoleRequirer:
		return RoleProvider
	case RolePeer:
		return RolePeer
	}
	panic(fmt.Errorf("unknown RelationRole: %q", r))
}

// RelationEndpoint represents one endpoint of a relation.
type RelationEndpoint struct {
	ServiceName   string
	Interface     string
	RelationName  string
	RelationRole  RelationRole
	RelationScope charm.RelationScope
}

// CanRelateTo returns whether a relation may be established between e and other.
func (e *RelationEndpoint) CanRelateTo(other *RelationEndpoint) bool {
	if e.Interface != other.Interface {
		return false
	}
	if e.RelationRole == RolePeer {
		// Peer relations do not currently work with multiple endpoints.
		return false
	}
	return e.RelationRole.counterpartRole() == other.RelationRole
}

// String returns the unique identifier of the relation endpoint.
func (e RelationEndpoint) String() string {
	return e.ServiceName + ":" + e.RelationName
}

// relationKey returns a string describing the relation defined by
// endpoints, for use in various contexts (including error messages).
func relationKey(endpoints []RelationEndpoint) string {
	names := []string{}
	for _, ep := range endpoints {
		names = append(names, ep.String())
	}
	sort.Strings(names)
	return strings.Join(names, " ")
}

// relationDoc is the internal representation of a Relation in MongoDB.
type relationDoc struct {
	Key       string `bson:"_id"`
	Id        int
	Endpoints []RelationEndpoint
	Life      Life
}

// Relation represents a relation between one or two service endpoints.
type Relation struct {
	st  *State
	doc relationDoc
}

func newRelation(st *State, doc *relationDoc) *Relation {
	return &Relation{
		st:  st,
		doc: *doc,
	}
}

func (r *Relation) String() string {
	return r.doc.Key
}

func (r *Relation) Refresh() error {
	doc := relationDoc{}
	err := r.st.relations.FindId(r.doc.Key).One(&doc)
	if err != nil {
		return fmt.Errorf("cannot refresh relation %v: %v", r, err)
	}
	r.doc = doc
	return nil
}

func (r *Relation) Life() Life {
	return r.doc.Life
}

// ensureLife changes the lifecycle state of the relation.
// See the Life type for more details.
func (r *Relation) ensureLife(life Life) error {
	if life == Alive {
		panic("cannot set life to alive")
	}
	sel := bson.D{
		{"_id", r.doc.Key},
		// $lte is used so that we don't overwrite a previous
		// change we don't know about.
		{"life", bson.D{{"$lte", life}}},
	}
	change := bson.D{{"$set", bson.D{{"life", life}}}}
	err := r.st.relations.Update(sel, change)
	if err == mgo.ErrNotFound {
		return nil
	}
	if err != nil {
		return fmt.Errorf("cannot set life to %v for relation %v: %v", life, r, err)
	}
	r.doc.Life = life
	return nil
}

// Kill sets the relation lifecycle to Dying if it is Alive.
// It does nothing otherwise.
func (r *Relation) Kill() error {
	return r.ensureLife(Dying)
}

// Die sets the relation lifecycle to Dead if it is Alive or Dying.
// It does nothing otherwise.
func (r *Relation) Die() error {
	return r.ensureLife(Dead)
}

// Id returns the integer internal relation key. This is exposed
// because the unit agent needs to expose a value derived from this
// (as JUJU_RELATION_ID) to allow relation hooks to differentiate
// between relations with different services.
func (r *Relation) Id() int {
	return r.doc.Id
}

// Endpoint returns the endpoint of the relation for the named service.
// If the service is not part of the relation, an error will be returned.
func (r *Relation) Endpoint(serviceName string) (RelationEndpoint, error) {
	for _, ep := range r.doc.Endpoints {
		if ep.ServiceName == serviceName {
			return ep, nil
		}
	}
	return RelationEndpoint{}, fmt.Errorf("service %q is not a member of %q", serviceName, r)
}

// RelatedEndpoints returns the endpoints of the relation r with which
// units of the named service will establish relations. If the service
// is not part of the relation r, an error will be returned.
func (r *Relation) RelatedEndpoints(serviceName string) ([]RelationEndpoint, error) {
	local, err := r.Endpoint(serviceName)
	if err != nil {
		return nil, err
	}
	role := local.RelationRole.counterpartRole()
	var eps []RelationEndpoint
	for _, ep := range r.doc.Endpoints {
		if ep.RelationRole == role {
			eps = append(eps, ep)
		}
	}
	if eps == nil {
		return nil, fmt.Errorf("no endpoints of %q relate to service %q", r, serviceName)
	}
	return eps, nil
}

// RelationUnit holds information about a single unit in a relation, and
// allows clients to conveniently access unit-specific functionality.
type RelationUnit struct {
	st       *State
	relation *Relation
	unit     *Unit
	endpoint RelationEndpoint
}

// Relation returns the relation associated with the unit.
func (ru *RelationUnit) Relation() *Relation {
	return ru.relation
}

// Endpoint returns the relation endpoint that defines the unit's
// participation in the relation.
func (ru *RelationUnit) Endpoint() RelationEndpoint {
	return ru.endpoint
}

// Unit returns a RelationUnit for the supplied unit.
func (r *Relation) Unit(u *Unit) (*RelationUnit, error) {
	ep, err := r.Endpoint(u.doc.Service)
	if err != nil {
		return nil, err
	}
	return &RelationUnit{
		st:       r.st,
		relation: r,
		unit:     u,
		endpoint: ep,
	}, nil
}
<commit_msg>mstate: relation Die/Kill update the cache<commit_after>package mstate

import (
	"fmt"
	"launchpad.net/juju-core/charm"
	"sort"
	"strings"
)

// RelationRole defines the role of a relation endpoint.
type RelationRole string

const (
	RoleProvider RelationRole = "provider"
	RoleRequirer RelationRole = "requirer"
	RolePeer     RelationRole = "peer"
)

// counterpartRole returns the RelationRole that this RelationRole
// can relate to.
// This should remain an internal method because the relation
// model does not guarantee that for every role there will
// necessarily exist a single counterpart role that is sensible
// for basing algorithms upon.
func (r RelationRole) counterpartRole() RelationRole {
	switch r {
	case RoleProvider:
		return RoleRequirer
	case RoleRequirer:
		return RoleProvider
	case RolePeer:
		return RolePeer
	}
	panic(fmt.Errorf("unknown RelationRole: %q", r))
}

// RelationEndpoint represents one endpoint of a relation.
type RelationEndpoint struct {
	ServiceName   string
	Interface     string
	RelationName  string
	RelationRole  RelationRole
	RelationScope charm.RelationScope
}

// CanRelateTo returns whether a relation may be established between e and other.
func (e *RelationEndpoint) CanRelateTo(other *RelationEndpoint) bool {
	if e.Interface != other.Interface {
		return false
	}
	if e.RelationRole == RolePeer {
		// Peer relations do not currently work with multiple endpoints.
		return false
	}
	return e.RelationRole.counterpartRole() == other.RelationRole
}

// String returns the unique identifier of the relation endpoint.
func (e RelationEndpoint) String() string {
	return e.ServiceName + ":" + e.RelationName
}

// relationKey returns a string describing the relation defined by
// endpoints, for use in various contexts (including error messages).
func relationKey(endpoints []RelationEndpoint) string {
	names := []string{}
	for _, ep := range endpoints {
		names = append(names, ep.String())
	}
	sort.Strings(names)
	return strings.Join(names, " ")
}

// relationDoc is the internal representation of a Relation in MongoDB.
type relationDoc struct {
	Key       string `bson:"_id"`
	Id        int
	Endpoints []RelationEndpoint
	Life      Life
}

// Relation represents a relation between one or two service endpoints.
type Relation struct {
	st  *State
	doc relationDoc
}

func newRelation(st *State, doc *relationDoc) *Relation {
	return &Relation{
		st:  st,
		doc: *doc,
	}
}

func (r *Relation) String() string {
	return r.doc.Key
}

func (r *Relation) Refresh() error {
	doc := relationDoc{}
	err := r.st.relations.FindId(r.doc.Key).One(&doc)
	if err != nil {
		return fmt.Errorf("cannot refresh relation %v: %v", r, err)
	}
	r.doc = doc
	return nil
}

func (r *Relation) Life() Life {
	return r.doc.Life
}

// Kill sets the relation lifecycle to Dying if it is Alive.
// It does nothing otherwise.
func (r *Relation) Kill() error {
	err := ensureLife(r.doc.Key, r.st, r.st.relations, "relation", Dying)
	if err != nil {
		return err
	}
	r.doc.Life = Dying
	return nil
}

// Die sets the relation lifecycle to Dead if it is Alive or Dying.
// It does nothing otherwise.
func (r *Relation) Die() error {
	err := ensureLife(r.doc.Key, r.st, r.st.relations, "relation", Dead)
	if err != nil {
		return err
	}
	r.doc.Life = Dead
	return nil
}

// Id returns the integer internal relation key. This is exposed
// because the unit agent needs to expose a value derived from this
// (as JUJU_RELATION_ID) to allow relation hooks to differentiate
// between relations with different services.
func (r *Relation) Id() int {
	return r.doc.Id
}

// Endpoint returns the endpoint of the relation for the named service.
// If the service is not part of the relation, an error will be returned.
func (r *Relation) Endpoint(serviceName string) (RelationEndpoint, error) {
	for _, ep := range r.doc.Endpoints {
		if ep.ServiceName == serviceName {
			return ep, nil
		}
	}
	return RelationEndpoint{}, fmt.Errorf("service %q is not a member of %q", serviceName, r)
}

// RelatedEndpoints returns the endpoints of the relation r with which
// units of the named service will establish relations. If the service
// is not part of the relation r, an error will be returned.
func (r *Relation) RelatedEndpoints(serviceName string) ([]RelationEndpoint, error) {
	local, err := r.Endpoint(serviceName)
	if err != nil {
		return nil, err
	}
	role := local.RelationRole.counterpartRole()
	var eps []RelationEndpoint
	for _, ep := range r.doc.Endpoints {
		if ep.RelationRole == role {
			eps = append(eps, ep)
		}
	}
	if eps == nil {
		return nil, fmt.Errorf("no endpoints of %q relate to service %q", r, serviceName)
	}
	return eps, nil
}

// RelationUnit holds information about a single unit in a relation, and
// allows clients to conveniently access unit-specific functionality.
type RelationUnit struct {
	st       *State
	relation *Relation
	unit     *Unit
	endpoint RelationEndpoint
}

// Relation returns the relation associated with the unit.
func (ru *RelationUnit) Relation() *Relation {
	return ru.relation
}

// Endpoint returns the relation endpoint that defines the unit's
// participation in the relation.
func (ru *RelationUnit) Endpoint() RelationEndpoint {
	return ru.endpoint
}

// Unit returns a RelationUnit for the supplied unit.
func (r *Relation) Unit(u *Unit) (*RelationUnit, error) {
	ep, err := r.Endpoint(u.doc.Service)
	if err != nil {
		return nil, err
	}
	return &RelationUnit{
		st:       r.st,
		relation: r,
		unit:     u,
		endpoint: ep,
	}, nil
}
<|endoftext|>


<commit_before>package handlers

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"

	log "github.com/Sirupsen/logrus"
	"github.com/stellar/gateway/protocols"
	"github.com/stellar/gateway/server"
)

// TransactionStatusResponse represents a response from the tx_status endpoint
type TransactionStatusResponse struct {
	protocols.SuccessResponse
	Status   string `json:"status"`
	RecvCode string `json:"recv_code,omitempty"`
	RefundTx string `json:"refund_tx,omitempty"`
	Msg      string `json:"msg,omitempty"`
}

// HandlerTxStatus implements /tx_status endpoint
func (rh *RequestHandler) HandlerTxStatus(w http.ResponseWriter, r *http.Request) {

	txid := r.URL.Query().Get("id")
	if txid == "" {
		log.Info("unable to get query parameter")
		server.Write(w, protocols.MissingParameterError)
		return
	}
	response := TransactionStatusResponse{}

	if rh.Config.Callbacks.TxStatus == "" {
		response.Status = "unknown"
	} else {
		endpoint := fmt.Sprintf(
			"%s?id=%s",
			rh.Config.Callbacks.TxStatus,
			txid,
		)

		_, err := url.Parse(endpoint)
		if err != nil {
			log.Error(err, "failed to parse tx status endpoint")
			server.Write(w, protocols.InternalServerError)
			return
		}
		resp, err := rh.Client.Get(endpoint)
		if err != nil {
			log.WithFields(log.Fields{
				"tx_status": rh.Config.Callbacks.TxStatus,
				"err":       err,
			}).Error("Error sending request to tx_status server")
			server.Write(w, protocols.InternalServerError)
			return
		}

		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Error("Error reading tx_status server response")
			server.Write(w, protocols.InternalServerError)
			return
		}

		switch resp.StatusCode {
		case http.StatusOK:
			err := json.Unmarshal(body, &response)
			if err != nil {
				log.WithFields(log.Fields{
					"tx_status": rh.Config.Callbacks.TxStatus,
					"body":      string(body),
				}).Error("Unable to decode tx_status response")
				server.Write(w, protocols.InternalServerError)
				return
			}
			if response.Status == "" {
				response.Status = "unknown"
			}

		default:
			response.Status = "unknown"
		}
	}

	w.Header().Set("Content-Type", "application/json")
	err := json.NewEncoder(w).Encode(response)
	if err != nil {
		log.Error("Error encoding tx status response")
		server.Write(w, protocols.InternalServerError)
		return
	}
}
<commit_msg>fix tx_status changes<commit_after>package handlers

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"

	log "github.com/Sirupsen/logrus"
	"github.com/stellar/gateway/protocols"
	"github.com/stellar/gateway/server"
	"github.com/stellar/go/protocols/compliance"
)

// HandlerTxStatus implements /tx_status endpoint
func (rh *RequestHandler) HandlerTxStatus(w http.ResponseWriter, r *http.Request) {

	txid := r.URL.Query().Get("id")
	if txid == "" {
		log.Info("unable to get query parameter")
		server.Write(w, protocols.NewMissingParameter("id"))
		return
	}
	response := compliance.TransactionStatusResponse{}

	if rh.Config.Callbacks.TxStatus == "" {
		response.Status = compliance.TransactionStatusUnknown
	} else {
		endpoint := fmt.Sprintf(
			"%s?id=%s",
			rh.Config.Callbacks.TxStatus,
			txid,
		)

		_, err := url.Parse(endpoint)
		if err != nil {
			log.Error(err, "failed to parse tx status endpoint")
			server.Write(w, protocols.InternalServerError)
			return
		}
		resp, err := rh.Client.Get(endpoint)
		if err != nil {
			log.WithFields(log.Fields{
				"tx_status": rh.Config.Callbacks.TxStatus,
				"err":       err,
			}).Error("Error sending request to tx_status server")
			server.Write(w, protocols.InternalServerError)
			return
		}

		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Error("Error reading tx_status server response")
			server.Write(w, protocols.InternalServerError)
			return
		}

		switch resp.StatusCode {
		case http.StatusOK:
			err := json.Unmarshal(body, &response)
			if err != nil {
				log.WithFields(log.Fields{
					"tx_status": rh.Config.Callbacks.TxStatus,
					"body":      string(body),
				}).Error("Unable to decode tx_status response")
				server.Write(w, protocols.InternalServerError)
				return
			}
			if response.Status == "" {
				response.Status = compliance.TransactionStatusUnknown
			}

		default:
			response.Status = compliance.TransactionStatusUnknown
		}
	}

	w.Header().Set("Content-Type", "application/json")
	err := json.NewEncoder(w).Encode(response)
	if err != nil {
		log.Error("Error encoding tx status response")
		server.Write(w, protocols.InternalServerError)
		return
	}
}
<|endoftext|>


<commit_before>package tsdb

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"sync"

	"github.com/influxdb/influxdb/influxql"
)

func NewStore(path string) *Store {
	return &Store{
		path:   path,
		Logger: log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

var (
	ErrShardNotFound = fmt.Errorf("shard not found")
)

type Store struct {
	mu   sync.RWMutex
	path string

	databaseIndexes map[string]*DatabaseIndex
	shards          map[uint64]*Shard

	Logger *log.Logger
}

func (s *Store) CreateShard(database, retentionPolicy string, shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// shard already exists
	if _, ok := s.shards[shardID]; ok {
		return nil
	}

	// created the db and retention policy dirs if they don't exist
	if err := os.MkdirAll(filepath.Join(s.path, database, retentionPolicy), 0700); err != nil {
		return err
	}

	// create the database index if it does not exist
	db, ok := s.databaseIndexes[database]
	if !ok {
		db = NewDatabaseIndex()
		s.databaseIndexes[database] = db
	}

	shardPath := filepath.Join(s.path, database, retentionPolicy, strconv.FormatUint(shardID, 10))
	shard := NewShard(db, shardPath)
	if err := shard.Open(); err != nil {
		return err
	}

	s.shards[shardID] = shard

	return nil
}

func (s *Store) Shard(shardID uint64) *Shard {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.shards[shardID]
}

func (s *Store) ValidateAggregateFieldsInStatement(shardID uint64, measurementName string, stmt *influxql.SelectStatement) error {
	s.mu.RLock()
	shard := s.shards[shardID]
	s.mu.RUnlock()
	if shard == nil {
		return ErrShardNotFound
	}
	return shard.ValidateAggregateFieldsInStatement(measurementName, stmt)
}

func (s *Store) DatabaseIndex(name string) *DatabaseIndex {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.databaseIndexes[name]
}

func (s *Store) Measurement(database, name string) *Measurement {
	s.mu.RLock()
	db := s.databaseIndexes[database]
	s.mu.RUnlock()
	if db == nil {
		return nil
	}
	return db.measurements[name]
}

// deleteSeries lopos through the local shards and deletes the series data and metadata for the passed in series keys
func (s *Store) deleteSeries(keys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.deleteSeries(keys); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) loadIndexes() error {
	dbs, err := ioutil.ReadDir(s.path)
	if err != nil {
		return err
	}
	for _, db := range dbs {
		if !db.IsDir() {
			s.Logger.Printf("Skipping database dir: %s. Not a directory", db.Name())
			continue
		}
		s.databaseIndexes[db.Name()] = NewDatabaseIndex()
	}
	return nil
}

func (s *Store) loadShards() error {
	// loop through the current database indexes
	for db := range s.databaseIndexes {
		rps, err := ioutil.ReadDir(filepath.Join(s.path, db))
		if err != nil {
			return err
		}

		for _, rp := range rps {
			// retention policies should be directories.  Skip anything that is not a dir.
			if !rp.IsDir() {
				s.Logger.Printf("Skipping retention policy dir: %s. Not a directory", rp.Name())
				continue
			}

			shards, err := ioutil.ReadDir(filepath.Join(s.path, db, rp.Name()))
			if err != nil {
				return err
			}
			for _, sh := range shards {
				path := filepath.Join(s.path, db, rp.Name(), sh.Name())

				// Shard file names are numeric shardIDs
				shardID, err := strconv.ParseUint(sh.Name(), 10, 64)
				if err != nil {
					s.Logger.Printf("Skipping shard: %s. Not a valid path", rp.Name())
					continue
				}

				shard := NewShard(s.databaseIndexes[db], path)
				shard.Open()
				s.shards[shardID] = shard
			}
		}
	}
	return nil

}

func (s *Store) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.shards = map[uint64]*Shard{}
	s.databaseIndexes = map[string]*DatabaseIndex{}

	// Create directory.
	if err := os.MkdirAll(s.path, 0777); err != nil {
		return err
	}

	// TODO: Start AE for Node
	if err := s.loadIndexes(); err != nil {
		return err
	}

	if err := s.loadShards(); err != nil {
		return err
	}

	return nil
}

func (s *Store) WriteToShard(shardID uint64, points []Point) error {
	sh, ok := s.shards[shardID]
	if !ok {
		return ErrShardNotFound
	}

	return sh.WritePoints(points)
}

func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, sh := range s.shards {
		if err := sh.Close(); err != nil {
			return err
		}
	}
	s.shards = nil
	s.databaseIndexes = nil

	return nil
}
<commit_msg>fix data race in WriteToShard<commit_after>package tsdb

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"sync"

	"github.com/influxdb/influxdb/influxql"
)

func NewStore(path string) *Store {
	return &Store{
		path:   path,
		Logger: log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

var (
	ErrShardNotFound = fmt.Errorf("shard not found")
)

type Store struct {
	mu   sync.RWMutex
	path string

	databaseIndexes map[string]*DatabaseIndex
	shards          map[uint64]*Shard

	Logger *log.Logger
}

func (s *Store) CreateShard(database, retentionPolicy string, shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// shard already exists
	if _, ok := s.shards[shardID]; ok {
		return nil
	}

	// created the db and retention policy dirs if they don't exist
	if err := os.MkdirAll(filepath.Join(s.path, database, retentionPolicy), 0700); err != nil {
		return err
	}

	// create the database index if it does not exist
	db, ok := s.databaseIndexes[database]
	if !ok {
		db = NewDatabaseIndex()
		s.databaseIndexes[database] = db
	}

	shardPath := filepath.Join(s.path, database, retentionPolicy, strconv.FormatUint(shardID, 10))
	shard := NewShard(db, shardPath)
	if err := shard.Open(); err != nil {
		return err
	}

	s.shards[shardID] = shard

	return nil
}

func (s *Store) Shard(shardID uint64) *Shard {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.shards[shardID]
}

func (s *Store) ValidateAggregateFieldsInStatement(shardID uint64, measurementName string, stmt *influxql.SelectStatement) error {
	s.mu.RLock()
	shard := s.shards[shardID]
	s.mu.RUnlock()
	if shard == nil {
		return ErrShardNotFound
	}
	return shard.ValidateAggregateFieldsInStatement(measurementName, stmt)
}

func (s *Store) DatabaseIndex(name string) *DatabaseIndex {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.databaseIndexes[name]
}

func (s *Store) Measurement(database, name string) *Measurement {
	s.mu.RLock()
	db := s.databaseIndexes[database]
	s.mu.RUnlock()
	if db == nil {
		return nil
	}
	return db.measurements[name]
}

// deleteSeries lopos through the local shards and deletes the series data and metadata for the passed in series keys
func (s *Store) deleteSeries(keys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.deleteSeries(keys); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) loadIndexes() error {
	dbs, err := ioutil.ReadDir(s.path)
	if err != nil {
		return err
	}
	for _, db := range dbs {
		if !db.IsDir() {
			s.Logger.Printf("Skipping database dir: %s. Not a directory", db.Name())
			continue
		}
		s.databaseIndexes[db.Name()] = NewDatabaseIndex()
	}
	return nil
}

func (s *Store) loadShards() error {
	// loop through the current database indexes
	for db := range s.databaseIndexes {
		rps, err := ioutil.ReadDir(filepath.Join(s.path, db))
		if err != nil {
			return err
		}

		for _, rp := range rps {
			// retention policies should be directories.  Skip anything that is not a dir.
			if !rp.IsDir() {
				s.Logger.Printf("Skipping retention policy dir: %s. Not a directory", rp.Name())
				continue
			}

			shards, err := ioutil.ReadDir(filepath.Join(s.path, db, rp.Name()))
			if err != nil {
				return err
			}
			for _, sh := range shards {
				path := filepath.Join(s.path, db, rp.Name(), sh.Name())

				// Shard file names are numeric shardIDs
				shardID, err := strconv.ParseUint(sh.Name(), 10, 64)
				if err != nil {
					s.Logger.Printf("Skipping shard: %s. Not a valid path", rp.Name())
					continue
				}

				shard := NewShard(s.databaseIndexes[db], path)
				shard.Open()
				s.shards[shardID] = shard
			}
		}
	}
	return nil

}

func (s *Store) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.shards = map[uint64]*Shard{}
	s.databaseIndexes = map[string]*DatabaseIndex{}

	// Create directory.
	if err := os.MkdirAll(s.path, 0777); err != nil {
		return err
	}

	// TODO: Start AE for Node
	if err := s.loadIndexes(); err != nil {
		return err
	}

	if err := s.loadShards(); err != nil {
		return err
	}

	return nil
}

func (s *Store) WriteToShard(shardID uint64, points []Point) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	sh, ok := s.shards[shardID]
	if !ok {
		return ErrShardNotFound
	}

	return sh.WritePoints(points)
}

func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, sh := range s.shards {
		if err := sh.Close(); err != nil {
			return err
		}
	}
	s.shards = nil
	s.databaseIndexes = nil

	return nil
}
<|endoftext|>