<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/matrix-org/dendrite/internal/caching"
	"github.com/matrix-org/dendrite/roomserver/state"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/dendrite/setup"
	"github.com/matrix-org/dendrite/setup/base"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
)

// This is a utility for inspecting state snapshots and running state resolution
// against real snapshots in an actual database.
// It takes one or more state snapshot NIDs as arguments, along with a room version
// to use for unmarshalling events, and will produce resolved output.
//
// Usage: ./resolve-state --roomversion=version snapshot [snapshot ...]
//   e.g. ./resolve-state --roomversion=5 1254 1235 1282

var roomVersion = flag.String("roomversion", "5", "the room version to parse events as")
var filterType = flag.String("filtertype", "", "the event types to filter on")
var difference = flag.Bool("difference", false, "whether to calculate the difference between snapshots")

// nolint:gocyclo
func main() {
	ctx := context.Background()
	cfg := setup.ParseFlags(true)
	cfg.Logging = append(cfg.Logging[:0], config.LogrusHook{
		Type:  "std",
		Level: "error",
	})
	cfg.ClientAPI.RegistrationDisabled = true
	base := base.NewBaseDendrite(cfg, "ResolveState", base.DisableMetrics)
	args := flag.Args()

	fmt.Println("Room version", *roomVersion)

	snapshotNIDs := []types.StateSnapshotNID{}
	for _, arg := range args {
		if i, err := strconv.Atoi(arg); err == nil {
			snapshotNIDs = append(snapshotNIDs, types.StateSnapshotNID(i))
		}
	}

	fmt.Println("Fetching", len(snapshotNIDs), "snapshot NIDs")

	roomserverDB, err := storage.Open(
		base, &cfg.RoomServer.Database,
		caching.NewRistrettoCache(128*1024*1024, time.Hour, true),
	)
	if err != nil {
		panic(err)
	}

	stateres := state.NewStateResolution(roomserverDB, &types.RoomInfo{
		RoomVersion: gomatrixserverlib.RoomVersion(*roomVersion),
	})

	if *difference {
		if len(snapshotNIDs) != 2 {
			panic("need exactly two state snapshot NIDs to calculate difference")
		}
		var removed, added []types.StateEntry
		removed, added, err = stateres.DifferenceBetweeenStateSnapshots(ctx, snapshotNIDs[0], snapshotNIDs[1])
		if err != nil {
			panic(err)
		}

		var eventNIDs []types.EventNID
		for _, entry := range append(removed, added...) {
			eventNIDs = append(eventNIDs, entry.EventNID)
		}

		var eventEntries []types.Event
		eventEntries, err = roomserverDB.Events(ctx, eventNIDs)
		if err != nil {
			panic(err)
		}

		events := make(map[types.EventNID]*gomatrixserverlib.Event, len(eventEntries))
		for _, entry := range eventEntries {
			events[entry.EventNID] = entry.Event
		}

		if len(removed) > 0 {
			fmt.Println("Removed:")
			for _, r := range removed {
				event := events[r.EventNID]
				fmt.Println()
				fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
				fmt.Printf("  %s\n", string(event.Content()))
			}
		}

		if len(removed) > 0 && len(added) > 0 {
			fmt.Println()
		}

		if len(added) > 0 {
			fmt.Println("Added:")
			for _, a := range added {
				event := events[a.EventNID]
				fmt.Println()
				fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
				fmt.Printf("  %s\n", string(event.Content()))
			}
		}

		return
	}

	var stateEntries []types.StateEntry
	for _, snapshotNID := range snapshotNIDs {
		var entries []types.StateEntry
		entries, err = stateres.LoadStateAtSnapshot(ctx, snapshotNID)
		if err != nil {
			panic(err)
		}
		stateEntries = append(stateEntries, entries...)
	}

	var eventNIDs []types.EventNID
	for _, entry := range stateEntries {
		eventNIDs = append(eventNIDs, entry.EventNID)
	}

	fmt.Println("Fetching", len(eventNIDs), "state events")
	eventEntries, err := roomserverDB.Events(ctx, eventNIDs)
	if err != nil {
		panic(err)
	}

	authEventIDMap := make(map[string]struct{})
	events := make([]*gomatrixserverlib.Event, len(eventEntries))
	for i := range eventEntries {
		events[i] = eventEntries[i].Event
		for _, authEventID := range eventEntries[i].AuthEventIDs() {
			authEventIDMap[authEventID] = struct{}{}
		}
	}

	authEventIDs := make([]string, 0, len(authEventIDMap))
	for authEventID := range authEventIDMap {
		authEventIDs = append(authEventIDs, authEventID)
	}

	fmt.Println("Fetching", len(authEventIDs), "auth events")
	authEventEntries, err := roomserverDB.EventsFromIDs(ctx, authEventIDs)
	if err != nil {
		panic(err)
	}

	authEvents := make([]*gomatrixserverlib.Event, len(authEventEntries))
	for i := range authEventEntries {
		authEvents[i] = authEventEntries[i].Event
	}

	fmt.Println("Resolving state")
	var resolved Events
	resolved, err = gomatrixserverlib.ResolveConflicts(
		gomatrixserverlib.RoomVersion(*roomVersion),
		events,
		authEvents,
	)
	if err != nil {
		panic(err)
	}

	fmt.Println("Resolved state contains", len(resolved), "events")
	sort.Sort(resolved)
	filteringEventType := *filterType
	count := 0
	for _, event := range resolved {
		if filteringEventType != "" && event.Type() != filteringEventType {
			continue
		}
		count++
		fmt.Println()
		fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
		fmt.Printf("  %s\n", string(event.Content()))
	}

	fmt.Println()
	fmt.Println("Returned", count, "state events after filtering")
}

type Events []*gomatrixserverlib.Event

func (e Events) Len() int {
	return len(e)
}

func (e Events) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

func (e Events) Less(i, j int) bool {
	typeDelta := strings.Compare(e[i].Type(), e[j].Type())
	if typeDelta < 0 {
		return true
	}
	if typeDelta > 0 {
		return false
	}
	stateKeyDelta := strings.Compare(*e[i].StateKey(), *e[j].StateKey())
	return stateKeyDelta < 0
}
<commit_msg>Optimise `resolve-state` tool<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/matrix-org/dendrite/internal/caching"
	"github.com/matrix-org/dendrite/roomserver/state"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/dendrite/setup"
	"github.com/matrix-org/dendrite/setup/base"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
)

// This is a utility for inspecting state snapshots and running state resolution
// against real snapshots in an actual database.
// It takes one or more state snapshot NIDs as arguments, along with a room version
// to use for unmarshalling events, and will produce resolved output.
//
// Usage: ./resolve-state --roomversion=version snapshot [snapshot ...]
//   e.g. ./resolve-state --roomversion=5 1254 1235 1282

var roomVersion = flag.String("roomversion", "5", "the room version to parse events as")
var filterType = flag.String("filtertype", "", "the event types to filter on")
var difference = flag.Bool("difference", false, "whether to calculate the difference between snapshots")

// nolint:gocyclo
func main() {
	ctx := context.Background()
	cfg := setup.ParseFlags(true)
	cfg.Logging = append(cfg.Logging[:0], config.LogrusHook{
		Type:  "std",
		Level: "error",
	})
	cfg.ClientAPI.RegistrationDisabled = true
	base := base.NewBaseDendrite(cfg, "ResolveState", base.DisableMetrics)
	args := flag.Args()

	fmt.Println("Room version", *roomVersion)

	snapshotNIDs := []types.StateSnapshotNID{}
	for _, arg := range args {
		if i, err := strconv.Atoi(arg); err == nil {
			snapshotNIDs = append(snapshotNIDs, types.StateSnapshotNID(i))
		}
	}

	fmt.Println("Fetching", len(snapshotNIDs), "snapshot NIDs")

	roomserverDB, err := storage.Open(
		base, &cfg.RoomServer.Database,
		caching.NewRistrettoCache(128*1024*1024, time.Hour, true),
	)
	if err != nil {
		panic(err)
	}

	stateres := state.NewStateResolution(roomserverDB, &types.RoomInfo{
		RoomVersion: gomatrixserverlib.RoomVersion(*roomVersion),
	})

	if *difference {
		if len(snapshotNIDs) != 2 {
			panic("need exactly two state snapshot NIDs to calculate difference")
		}
		var removed, added []types.StateEntry
		removed, added, err = stateres.DifferenceBetweeenStateSnapshots(ctx, snapshotNIDs[0], snapshotNIDs[1])
		if err != nil {
			panic(err)
		}

		eventNIDMap := map[types.EventNID]struct{}{}
		for _, entry := range append(removed, added...) {
			eventNIDMap[entry.EventNID] = struct{}{}
		}

		eventNIDs := make([]types.EventNID, 0, len(eventNIDMap))
		for eventNID := range eventNIDMap {
			eventNIDs = append(eventNIDs, eventNID)
		}

		var eventEntries []types.Event
		eventEntries, err = roomserverDB.Events(ctx, eventNIDs)
		if err != nil {
			panic(err)
		}

		events := make(map[types.EventNID]*gomatrixserverlib.Event, len(eventEntries))
		for _, entry := range eventEntries {
			events[entry.EventNID] = entry.Event
		}

		if len(removed) > 0 {
			fmt.Println("Removed:")
			for _, r := range removed {
				event := events[r.EventNID]
				fmt.Println()
				fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
				fmt.Printf("  %s\n", string(event.Content()))
			}
		}

		if len(removed) > 0 && len(added) > 0 {
			fmt.Println()
		}

		if len(added) > 0 {
			fmt.Println("Added:")
			for _, a := range added {
				event := events[a.EventNID]
				fmt.Println()
				fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
				fmt.Printf("  %s\n", string(event.Content()))
			}
		}

		return
	}

	var stateEntries []types.StateEntry
	for _, snapshotNID := range snapshotNIDs {
		var entries []types.StateEntry
		entries, err = stateres.LoadStateAtSnapshot(ctx, snapshotNID)
		if err != nil {
			panic(err)
		}
		stateEntries = append(stateEntries, entries...)
	}

	eventNIDMap := map[types.EventNID]struct{}{}
	for _, entry := range stateEntries {
		eventNIDMap[entry.EventNID] = struct{}{}
	}

	eventNIDs := make([]types.EventNID, 0, len(eventNIDMap))
	for eventNID := range eventNIDMap {
		eventNIDs = append(eventNIDs, eventNID)
	}

	fmt.Println("Fetching", len(eventNIDMap), "state events")
	eventEntries, err := roomserverDB.Events(ctx, eventNIDs)
	if err != nil {
		panic(err)
	}

	authEventIDMap := make(map[string]struct{})
	events := make([]*gomatrixserverlib.Event, len(eventEntries))
	for i := range eventEntries {
		events[i] = eventEntries[i].Event
		for _, authEventID := range eventEntries[i].AuthEventIDs() {
			authEventIDMap[authEventID] = struct{}{}
		}
	}

	authEventIDs := make([]string, 0, len(authEventIDMap))
	for authEventID := range authEventIDMap {
		authEventIDs = append(authEventIDs, authEventID)
	}

	fmt.Println("Fetching", len(authEventIDs), "auth events")
	authEventEntries, err := roomserverDB.EventsFromIDs(ctx, authEventIDs)
	if err != nil {
		panic(err)
	}

	authEvents := make([]*gomatrixserverlib.Event, len(authEventEntries))
	for i := range authEventEntries {
		authEvents[i] = authEventEntries[i].Event
	}

	fmt.Println("Resolving state")
	var resolved Events
	resolved, err = gomatrixserverlib.ResolveConflicts(
		gomatrixserverlib.RoomVersion(*roomVersion),
		events,
		authEvents,
	)
	if err != nil {
		panic(err)
	}

	fmt.Println("Resolved state contains", len(resolved), "events")
	sort.Sort(resolved)
	filteringEventType := *filterType
	count := 0
	for _, event := range resolved {
		if filteringEventType != "" && event.Type() != filteringEventType {
			continue
		}
		count++
		fmt.Println()
		fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
		fmt.Printf("  %s\n", string(event.Content()))
	}

	fmt.Println()
	fmt.Println("Returned", count, "state events after filtering")
}

type Events []*gomatrixserverlib.Event

func (e Events) Len() int {
	return len(e)
}

func (e Events) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

func (e Events) Less(i, j int) bool {
	typeDelta := strings.Compare(e[i].Type(), e[j].Type())
	if typeDelta < 0 {
		return true
	}
	if typeDelta > 0 {
		return false
	}
	stateKeyDelta := strings.Compare(*e[i].StateKey(), *e[j].StateKey())
	return stateKeyDelta < 0
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"strconv"

	"github.com/matrix-org/dendrite/internal/caching"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/dendrite/setup"
	"github.com/matrix-org/gomatrixserverlib"
)

// This is a utility for inspecting state snapshots and running state resolution
// against real snapshots in an actual database.
// It takes one or more state snapshot NIDs as arguments, along with a room version
// to use for unmarshalling events, and will produce resolved output.
//
// Usage: ./resolve-state --roomversion=version snapshot [snapshot ...]
//   e.g. ./resolve-state --roomversion=5 1254 1235 1282

var roomVersion = flag.String("roomversion", "5", "the room version to parse events as")

func main() {
	ctx := context.Background()
	cfg := setup.ParseFlags(true)
	args := os.Args[1:]

	fmt.Println("Room version", *roomVersion)

	snapshotNIDs := []types.StateSnapshotNID{}
	for _, arg := range args {
		if i, err := strconv.Atoi(arg); err == nil {
			snapshotNIDs = append(snapshotNIDs, types.StateSnapshotNID(i))
		}
	}

	fmt.Println("Fetching", len(snapshotNIDs), "snapshot NIDs")

	cache, err := caching.NewInMemoryLRUCache(true)
	if err != nil {
		panic(err)
	}

	roomserverDB, err := storage.Open(nil, &cfg.RoomServer.Database, cache)
	if err != nil {
		panic(err)
	}

	blockNIDs, err := roomserverDB.StateBlockNIDs(ctx, snapshotNIDs)
	if err != nil {
		panic(err)
	}

	var stateEntries []types.StateEntryList
	for _, list := range blockNIDs {
		entries, err2 := roomserverDB.StateEntries(ctx, list.StateBlockNIDs)
		if err2 != nil {
			panic(err2)
		}
		stateEntries = append(stateEntries, entries...)
	}

	var eventNIDs []types.EventNID
	for _, entry := range stateEntries {
		for _, e := range entry.StateEntries {
			eventNIDs = append(eventNIDs, e.EventNID)
		}
	}

	fmt.Println("Fetching", len(eventNIDs), "state events")
	eventEntries, err := roomserverDB.Events(ctx, eventNIDs)
	if err != nil {
		panic(err)
	}

	authEventIDMap := make(map[string]struct{})
	events := make([]*gomatrixserverlib.Event, len(eventEntries))
	for i := range eventEntries {
		events[i] = eventEntries[i].Event
		for _, authEventID := range eventEntries[i].AuthEventIDs() {
			authEventIDMap[authEventID] = struct{}{}
		}
	}

	authEventIDs := make([]string, 0, len(authEventIDMap))
	for authEventID := range authEventIDMap {
		authEventIDs = append(authEventIDs, authEventID)
	}

	fmt.Println("Fetching", len(authEventIDs), "auth events")
	authEventEntries, err := roomserverDB.EventsFromIDs(ctx, authEventIDs)
	if err != nil {
		panic(err)
	}

	authEvents := make([]*gomatrixserverlib.Event, len(authEventEntries))
	for i := range authEventEntries {
		authEvents[i] = authEventEntries[i].Event
	}

	fmt.Println("Resolving state")
	resolved, err := gomatrixserverlib.ResolveConflicts(
		gomatrixserverlib.RoomVersion(*roomVersion),
		events,
		authEvents,
	)
	if err != nil {
		panic(err)
	}

	fmt.Println("Resolved state contains", len(resolved), "events")
	for _, event := range resolved {
		fmt.Println()
		fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
		fmt.Printf("  %s\n", string(event.Content()))
	}
}
<commit_msg>Fix the `resolve-state` tool<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"strconv"

	"github.com/matrix-org/dendrite/internal/caching"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/dendrite/setup"
	"github.com/matrix-org/dendrite/setup/base"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
)

// This is a utility for inspecting state snapshots and running state resolution
// against real snapshots in an actual database.
// It takes one or more state snapshot NIDs as arguments, along with a room version
// to use for unmarshalling events, and will produce resolved output.
//
// Usage: ./resolve-state --roomversion=version snapshot [snapshot ...]
//   e.g. ./resolve-state --roomversion=5 1254 1235 1282

var roomVersion = flag.String("roomversion", "5", "the room version to parse events as")
var filterType = flag.String("filtertype", "", "the event types to filter on")

func main() {
	ctx := context.Background()
	cfg := setup.ParseFlags(true)
	cfg.Logging = append(cfg.Logging[:0], config.LogrusHook{
		Type:  "std",
		Level: "error",
	})
	base := base.NewBaseDendrite(cfg, "ResolveState", base.DisableMetrics)
	args := flag.Args()

	fmt.Println("Room version", *roomVersion)

	snapshotNIDs := []types.StateSnapshotNID{}
	for _, arg := range args {
		if i, err := strconv.Atoi(arg); err == nil {
			snapshotNIDs = append(snapshotNIDs, types.StateSnapshotNID(i))
		}
	}

	fmt.Println("Fetching", len(snapshotNIDs), "snapshot NIDs")

	cache, err := caching.NewInMemoryLRUCache(true)
	if err != nil {
		panic(err)
	}

	roomserverDB, err := storage.Open(base, &cfg.RoomServer.Database, cache)
	if err != nil {
		panic(err)
	}

	blockNIDs, err := roomserverDB.StateBlockNIDs(ctx, snapshotNIDs)
	if err != nil {
		panic(err)
	}

	var stateEntries []types.StateEntryList
	for _, list := range blockNIDs {
		entries, err2 := roomserverDB.StateEntries(ctx, list.StateBlockNIDs)
		if err2 != nil {
			panic(err2)
		}
		stateEntries = append(stateEntries, entries...)
	}

	var eventNIDs []types.EventNID
	for _, entry := range stateEntries {
		for _, e := range entry.StateEntries {
			eventNIDs = append(eventNIDs, e.EventNID)
		}
	}

	fmt.Println("Fetching", len(eventNIDs), "state events")
	eventEntries, err := roomserverDB.Events(ctx, eventNIDs)
	if err != nil {
		panic(err)
	}

	authEventIDMap := make(map[string]struct{})
	events := make([]*gomatrixserverlib.Event, len(eventEntries))
	for i := range eventEntries {
		events[i] = eventEntries[i].Event
		for _, authEventID := range eventEntries[i].AuthEventIDs() {
			authEventIDMap[authEventID] = struct{}{}
		}
	}

	authEventIDs := make([]string, 0, len(authEventIDMap))
	for authEventID := range authEventIDMap {
		authEventIDs = append(authEventIDs, authEventID)
	}

	fmt.Println("Fetching", len(authEventIDs), "auth events")
	authEventEntries, err := roomserverDB.EventsFromIDs(ctx, authEventIDs)
	if err != nil {
		panic(err)
	}

	authEvents := make([]*gomatrixserverlib.Event, len(authEventEntries))
	for i := range authEventEntries {
		authEvents[i] = authEventEntries[i].Event
	}

	fmt.Println("Resolving state")
	resolved, err := gomatrixserverlib.ResolveConflicts(
		gomatrixserverlib.RoomVersion(*roomVersion),
		events,
		authEvents,
	)
	if err != nil {
		panic(err)
	}

	fmt.Println("Resolved state contains", len(resolved), "events")
	filteringEventType := *filterType
	count := 0
	for _, event := range resolved {
		if filteringEventType != "" && event.Type() != filteringEventType {
			continue
		}
		count++
		fmt.Println()
		fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
		fmt.Printf("  %s\n", string(event.Content()))
	}

	fmt.Println()
	fmt.Println("Returned", count, "state events after filtering")
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"

	"github.com/spf13/cobra"
	"golang.org/x/sync/errgroup"

	"github.com/restic/restic/internal/backend"
	"github.com/restic/restic/internal/debug"
	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/repository"
	"github.com/restic/restic/internal/restic"
	"github.com/restic/restic/internal/walker"
)

var cmdRewrite = &cobra.Command{
	Use:   "rewrite [flags] [snapshotID ...]",
	Short: "Rewrite existing snapshots",
	Long: `
The "rewrite" command excludes files from existing snapshots.

By default 'rewrite' will create new snapshots that will contains same data as
the source snapshots but without excluded files. All metadata (time, host, tags)
will be preserved. The special tag 'rewrite' will be added to new snapshots to
distinguish it from the source (unless --inplace is used).

If --inplace option is used, old snapshot will be removed from repository.

Snapshots to rewrite are specified using --host, --tag, --path or by providing
a list of snapshot ids. Not specifying a snapshot id will rewrite all snapshots.

Please note, that this command only creates new snapshots. In order to delete
data from the repository use 'prune' command.

EXIT STATUS
===========

Exit status is 0 if the command was successful, and non-zero if there was any error.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runRewrite(cmd.Context(), rewriteOptions, globalOptions, args)
	},
}

// RewriteOptions collects all options for the rewrite command.
type RewriteOptions struct {
	Inplace bool
	DryRun  bool

	snapshotFilterOptions
	excludePatternOptions
}

var rewriteOptions RewriteOptions

func init() {
	cmdRoot.AddCommand(cmdRewrite)

	f := cmdRewrite.Flags()
	f.BoolVarP(&rewriteOptions.Inplace, "inplace", "", false, "replace existing snapshots")
	f.BoolVarP(&rewriteOptions.DryRun, "dry-run", "n", false, "do not do anything, just print what would be done")

	initMultiSnapshotFilterOptions(f, &rewriteOptions.snapshotFilterOptions, true)
	initExcludePatternOptions(f, &rewriteOptions.excludePatternOptions)
}

func rewriteSnapshot(ctx context.Context, repo *repository.Repository, sn *restic.Snapshot, opts RewriteOptions) (bool, error) {
	if sn.Tree == nil {
		return false, errors.Errorf("snapshot %v has nil tree", sn.ID().Str())
	}

	rejectByNameFuncs, err := collectExcludePatterns(opts.excludePatternOptions)
	if err != nil {
		return false, err
	}

	selectByName := func(nodepath string) bool {
		for _, reject := range rejectByNameFuncs {
			if reject(nodepath) {
				return false
			}
		}
		return true
	}

	wg, wgCtx := errgroup.WithContext(ctx)
	repo.StartPackUploader(wgCtx, wg)

	var filteredTree restic.ID
	wg.Go(func() error {
		filteredTree, err = walker.FilterTree(wgCtx, repo, "/", *sn.Tree, &walker.TreeFilterVisitor{
			SelectByName: selectByName,
			PrintExclude: func(path string) { Verbosef(fmt.Sprintf("excluding %s\n", path)) },
		})
		if err != nil {
			return err
		}

		return repo.Flush(wgCtx)
	})
	err = wg.Wait()
	if err != nil {
		return false, err
	}

	if filteredTree == *sn.Tree {
		debug.Log("Snapshot %v not modified", sn)
		return false, nil
	}

	debug.Log("Snapshot %v modified", sn)
	if opts.DryRun {
		Verbosef("would save new snapshot\n")
		return true, nil
	}

	// Retain the original snapshot id over all tag changes.
	if sn.Original == nil {
		sn.Original = sn.ID()
	}
	*sn.Tree = filteredTree

	if !opts.Inplace {
		sn.AddTags([]string{"rewrite"})
	}

	// Save the new snapshot.
	id, err := restic.SaveSnapshot(ctx, repo, sn)
	if err != nil {
		return false, err
	}

	if opts.Inplace {
		h := restic.Handle{Type: restic.SnapshotFile, Name: sn.ID().String()}
		if err = repo.Backend().Remove(ctx, h); err != nil {
			return false, err
		}

		debug.Log("old snapshot %v removed", sn.ID())
	}
	Verbosef("new snapshot saved as %v\n", id)
	return true, nil
}

func runRewrite(ctx context.Context, opts RewriteOptions, gopts GlobalOptions, args []string) error {

	if len(opts.ExcludeFiles) == 0 && len(opts.Excludes) == 0 && len(opts.InsensitiveExcludes) == 0 {
		return errors.Fatal("Nothing to do: no excludes provided")
	}

	repo, err := OpenRepository(ctx, gopts)
	if err != nil {
		return err
	}

	if !opts.DryRun {
		var lock *restic.Lock
		var err error
		if opts.Inplace {
			Verbosef("create exclusive lock for repository\n")
			lock, ctx, err = lockRepoExclusive(ctx, repo)
		} else {
			lock, ctx, err = lockRepo(ctx, repo)
		}
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	} else {
		repo.SetDryRun()
	}

	snapshotLister, err := backend.MemorizeList(ctx, repo.Backend(), restic.SnapshotFile)
	if err != nil {
		return err
	}

	if err = repo.LoadIndex(ctx); err != nil {
		return err
	}

	changedCount := 0
	for sn := range FindFilteredSnapshots(ctx, snapshotLister, repo, opts.Hosts, opts.Tags, opts.Paths, args) {
		Verbosef("\nsnapshot %s of %v at %s)\n", sn.ID().Str(), sn.Paths, sn.Time)
		changed, err := rewriteSnapshot(ctx, repo, sn, opts)
		if err != nil {
			return errors.Fatalf("unable to rewrite snapshot ID %q: %v", sn.ID().Str(), err)
		}
		if changed {
			changedCount++
		}
	}

	Verbosef("\n")
	if changedCount == 0 {
		Verbosef("no snapshots were modified\n")
	} else {
		if !opts.DryRun {
			Verbosef("modified %v snapshots\n", changedCount)
		} else {
			Verbosef("would modify %v snapshots\n", changedCount)
		}
	}

	return nil
}
<commit_msg>rewrite: rename --inplace to --forget<commit_after>package main

import (
	"context"
	"fmt"

	"github.com/spf13/cobra"
	"golang.org/x/sync/errgroup"

	"github.com/restic/restic/internal/backend"
	"github.com/restic/restic/internal/debug"
	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/repository"
	"github.com/restic/restic/internal/restic"
	"github.com/restic/restic/internal/walker"
)

var cmdRewrite = &cobra.Command{
	Use:   "rewrite [flags] [snapshotID ...]",
	Short: "Rewrite existing snapshots",
	Long: `
The "rewrite" command excludes files from existing snapshots.

By default 'rewrite' will create new snapshots that will contains same data as
the source snapshots but without excluded files. All metadata (time, host, tags)
will be preserved. The special tag 'rewrite' will be added to new snapshots to
distinguish it from the source (unless --forget is used).

If --forget option is used, old snapshot will be removed from repository.

Snapshots to rewrite are specified using --host, --tag, --path or by providing
a list of snapshot ids. Not specifying a snapshot id will rewrite all snapshots.

Please note, that this command only creates new snapshots. In order to delete
data from the repository use 'prune' command.

EXIT STATUS
===========

Exit status is 0 if the command was successful, and non-zero if there was any error.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runRewrite(cmd.Context(), rewriteOptions, globalOptions, args)
	},
}

// RewriteOptions collects all options for the rewrite command.
type RewriteOptions struct {
	Forget bool
	DryRun bool

	snapshotFilterOptions
	excludePatternOptions
}

var rewriteOptions RewriteOptions

func init() {
	cmdRoot.AddCommand(cmdRewrite)

	f := cmdRewrite.Flags()
	f.BoolVarP(&rewriteOptions.Forget, "forget", "", false, "replace existing snapshots")
	f.BoolVarP(&rewriteOptions.DryRun, "dry-run", "n", false, "do not do anything, just print what would be done")

	initMultiSnapshotFilterOptions(f, &rewriteOptions.snapshotFilterOptions, true)
	initExcludePatternOptions(f, &rewriteOptions.excludePatternOptions)
}

func rewriteSnapshot(ctx context.Context, repo *repository.Repository, sn *restic.Snapshot, opts RewriteOptions) (bool, error) {
	if sn.Tree == nil {
		return false, errors.Errorf("snapshot %v has nil tree", sn.ID().Str())
	}

	rejectByNameFuncs, err := collectExcludePatterns(opts.excludePatternOptions)
	if err != nil {
		return false, err
	}

	selectByName := func(nodepath string) bool {
		for _, reject := range rejectByNameFuncs {
			if reject(nodepath) {
				return false
			}
		}
		return true
	}

	wg, wgCtx := errgroup.WithContext(ctx)
	repo.StartPackUploader(wgCtx, wg)

	var filteredTree restic.ID
	wg.Go(func() error {
		filteredTree, err = walker.FilterTree(wgCtx, repo, "/", *sn.Tree, &walker.TreeFilterVisitor{
			SelectByName: selectByName,
			PrintExclude: func(path string) { Verbosef(fmt.Sprintf("excluding %s\n", path)) },
		})
		if err != nil {
			return err
		}

		return repo.Flush(wgCtx)
	})
	err = wg.Wait()
	if err != nil {
		return false, err
	}

	if filteredTree == *sn.Tree {
		debug.Log("Snapshot %v not modified", sn)
		return false, nil
	}

	debug.Log("Snapshot %v modified", sn)
	if opts.DryRun {
		Verbosef("would save new snapshot\n")
		return true, nil
	}

	// Retain the original snapshot id over all tag changes.
	if sn.Original == nil {
		sn.Original = sn.ID()
	}
	*sn.Tree = filteredTree

	if !opts.Forget {
		sn.AddTags([]string{"rewrite"})
	}

	// Save the new snapshot.
	id, err := restic.SaveSnapshot(ctx, repo, sn)
	if err != nil {
		return false, err
	}

	if opts.Forget {
		h := restic.Handle{Type: restic.SnapshotFile, Name: sn.ID().String()}
		if err = repo.Backend().Remove(ctx, h); err != nil {
			return false, err
		}

		debug.Log("old snapshot %v removed", sn.ID())
	}
	Verbosef("new snapshot saved as %v\n", id)
	return true, nil
}

func runRewrite(ctx context.Context, opts RewriteOptions, gopts GlobalOptions, args []string) error {

	if len(opts.ExcludeFiles) == 0 && len(opts.Excludes) == 0 && len(opts.InsensitiveExcludes) == 0 {
		return errors.Fatal("Nothing to do: no excludes provided")
	}

	repo, err := OpenRepository(ctx, gopts)
	if err != nil {
		return err
	}

	if !opts.DryRun {
		var lock *restic.Lock
		var err error
		if opts.Forget {
			Verbosef("create exclusive lock for repository\n")
			lock, ctx, err = lockRepoExclusive(ctx, repo)
		} else {
			lock, ctx, err = lockRepo(ctx, repo)
		}
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	} else {
		repo.SetDryRun()
	}

	snapshotLister, err := backend.MemorizeList(ctx, repo.Backend(), restic.SnapshotFile)
	if err != nil {
		return err
	}

	if err = repo.LoadIndex(ctx); err != nil {
		return err
	}

	changedCount := 0
	for sn := range FindFilteredSnapshots(ctx, snapshotLister, repo, opts.Hosts, opts.Tags, opts.Paths, args) {
		Verbosef("\nsnapshot %s of %v at %s)\n", sn.ID().Str(), sn.Paths, sn.Time)
		changed, err := rewriteSnapshot(ctx, repo, sn, opts)
		if err != nil {
			return errors.Fatalf("unable to rewrite snapshot ID %q: %v", sn.ID().Str(), err)
		}
		if changed {
			changedCount++
		}
	}

	Verbosef("\n")
	if changedCount == 0 {
		Verbosef("no snapshots were modified\n")
	} else {
		if !opts.DryRun {
			Verbosef("modified %v snapshots\n", changedCount)
		} else {
			Verbosef("would modify %v snapshots\n", changedCount)
		}
	}

	return nil
}
<|endoftext|>


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

// The rmplaysnippet binary removes a code snippet from play.golang.org given its URL
// or ID. It will always connect to the production datastore instance, ignoring any
// local value of DATASTORE_EMULATOR_HOST.
package main

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	"cloud.google.com/go/datastore"
	"golang.org/x/build/buildenv"
)

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s {http(s)://play.golang.org/p/<id> | <id>}\n", os.Args[0])
}

func main() {
	if len(os.Args) != 2 {
		usage()
		os.Exit(2)
	}

	snippetID := os.Args[1]
	prefixes := []string{"https://play.golang.org/p/", "http://play.golang.org/p/"}
	for _, p := range prefixes {
		if strings.HasPrefix(os.Args[1], p) {
			snippetID = strings.TrimPrefix(os.Args[1], p)
			break
		}
	}
	if snippetID == "" {
		usage()
		os.Exit(2)
	}
	if strings.Contains(snippetID, "/") {
		usage()
		fmt.Fprintf(os.Stderr, "Invalid Snippet ID %q (contains slash)\n", snippetID)
		os.Exit(2)
	}

	fmt.Printf("Really delete Snippet with ID %q? [y,N]: ", snippetID)
	var confirm string
	fmt.Scanln(&confirm)
	if !strings.HasPrefix(strings.ToLower(confirm), "y") {
		fmt.Printf("Aborting ...\n")
		os.Exit(0)
	}

	buildenv.CheckUserCredentials()

	// Don't attempt to connect to a locally-running datastore instance.
	if err := os.Setenv("DATASTORE_EMULATOR_HOST", ""); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to clear env var DATASTORE_EMULATOR_HOST: %v\n", err)
		os.Exit(1)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	client, err := datastore.NewClient(ctx, "golang-org")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create Datastore client: %v\n", err)
		os.Exit(1)
	}
	k := datastore.NameKey("Snippet", snippetID, nil)
	if client.Get(ctx, k, new(struct{})) == datastore.ErrNoSuchEntity {
		fmt.Fprintf(os.Stderr, "Snippet with ID %q does not exist\n", snippetID)
		os.Exit(0)
	}
	fmt.Printf("Deleting snippet %q ...\n", snippetID)
	if err := client.Delete(ctx, k); err != nil {
		fmt.Fprintf(os.Stderr, "Unable to delete Snippet with ID %q: %v\n", snippetID, err)
		fmt.Fprintf(os.Stderr, "rmplaysnippet requires Application Default Credentials.\n")
		fmt.Fprintf(os.Stderr, "Did you run `gcloud auth application-default login`?\n")
		os.Exit(1)
	}
	fmt.Printf("Snippet with ID %q deleted\n", snippetID)
}
<commit_msg>rmplaysnippet: recognize go.dev/play as playground prefix<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The rmplaysnippet binary removes a code snippet from play.golang.org given its URL
// or ID. It will always connect to the production datastore instance, ignoring any
// local value of DATASTORE_EMULATOR_HOST.
package main

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	"cloud.google.com/go/datastore"
	"golang.org/x/build/buildenv"
)

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s {http(s)://play.golang.org/p/<id> | <id>}\n", os.Args[0])
}

func main() {
	if len(os.Args) != 2 {
		usage()
		os.Exit(2)
	}

	snippetID := os.Args[1]
	prefixes := []string{
		"https://play.golang.org/p/",
		"http://play.golang.org/p/",
		"https://go.dev/play/p/",
		"http://go.dev/play/p/",
	}
	for _, p := range prefixes {
		if strings.HasPrefix(os.Args[1], p) {
			snippetID = strings.TrimPrefix(os.Args[1], p)
			break
		}
	}
	if snippetID == "" {
		usage()
		os.Exit(2)
	}
	if strings.Contains(snippetID, "/") {
		usage()
		fmt.Fprintf(os.Stderr, "Invalid Snippet ID %q (contains slash)\n", snippetID)
		os.Exit(2)
	}

	fmt.Printf("Really delete Snippet with ID %q? [y,N]: ", snippetID)
	var confirm string
	fmt.Scanln(&confirm)
	if !strings.HasPrefix(strings.ToLower(confirm), "y") {
		fmt.Printf("Aborting ...\n")
		os.Exit(0)
	}

	buildenv.CheckUserCredentials()

	// Don't attempt to connect to a locally-running datastore instance.
	if err := os.Setenv("DATASTORE_EMULATOR_HOST", ""); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to clear env var DATASTORE_EMULATOR_HOST: %v\n", err)
		os.Exit(1)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	client, err := datastore.NewClient(ctx, "golang-org")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create Datastore client: %v\n", err)
		os.Exit(1)
	}
	k := datastore.NameKey("Snippet", snippetID, nil)
	if client.Get(ctx, k, new(struct{})) == datastore.ErrNoSuchEntity {
		fmt.Fprintf(os.Stderr, "Snippet with ID %q does not exist\n", snippetID)
		os.Exit(0)
	}
	fmt.Printf("Deleting snippet %q ...\n", snippetID)
	if err := client.Delete(ctx, k); err != nil {
		fmt.Fprintf(os.Stderr, "Unable to delete Snippet with ID %q: %v\n", snippetID, err)
		fmt.Fprintf(os.Stderr, "rmplaysnippet requires Application Default Credentials.\n")
		fmt.Fprintf(os.Stderr, "Did you run `gcloud auth application-default login`?\n")
		os.Exit(1)
	}
	fmt.Printf("Snippet with ID %q deleted\n", snippetID)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	"github.com/minio/minio/internal/auth"
	xhttp "github.com/minio/minio/internal/http"
	"github.com/minio/minio/internal/logger"
	iampolicy "github.com/minio/pkg/iam/policy"
)

// http Header "x-amz-content-sha256" == "UNSIGNED-PAYLOAD" indicates that the
// client did not calculate sha256 of the payload.
const unsignedPayload = "UNSIGNED-PAYLOAD"

// skipContentSha256Cksum returns true if caller needs to skip
// payload checksum, false if not.
func skipContentSha256Cksum(r *http.Request) bool {
	var (
		v  []string
		ok bool
	)

	if isRequestPresignedSignatureV4(r) {
		v, ok = r.Form[xhttp.AmzContentSha256]
		if !ok {
			v, ok = r.Header[xhttp.AmzContentSha256]
		}
	} else {
		v, ok = r.Header[xhttp.AmzContentSha256]
	}

	// If x-amz-content-sha256 is set and the value is not
	// 'UNSIGNED-PAYLOAD' we should validate the content sha256.
	return !(ok && v[0] != unsignedPayload)
}

// Returns SHA256 for calculating canonical-request.
func getContentSha256Cksum(r *http.Request, stype serviceType) string {
	if stype == serviceSTS {
		payload, err := ioutil.ReadAll(io.LimitReader(r.Body, stsRequestBodyLimit))
		if err != nil {
			logger.CriticalIf(GlobalContext, err)
		}
		sum256 := sha256.Sum256(payload)
		r.Body = ioutil.NopCloser(bytes.NewReader(payload))
		return hex.EncodeToString(sum256[:])
	}

	var (
		defaultSha256Cksum string
		v                  []string
		ok                 bool
	)

	// For a presigned request we look at the query param for sha256.
	if isRequestPresignedSignatureV4(r) {
		// X-Amz-Content-Sha256, if not set in presigned requests, checksum
		// will default to 'UNSIGNED-PAYLOAD'.
		defaultSha256Cksum = unsignedPayload
		v, ok = r.Form[xhttp.AmzContentSha256]
		if !ok {
			v, ok = r.Header[xhttp.AmzContentSha256]
		}
	} else {
		// X-Amz-Content-Sha256, if not set in signed requests, checksum
		// will default to sha256([]byte("")).
		defaultSha256Cksum = emptySHA256
		v, ok = r.Header[xhttp.AmzContentSha256]
	}

	// We found 'X-Amz-Content-Sha256' return the captured value.
	if ok {
		return v[0]
	}

	// We couldn't find 'X-Amz-Content-Sha256'.
	return defaultSha256Cksum
}

// isValidRegion - verify if incoming region value is valid with configured Region.
func isValidRegion(reqRegion string, confRegion string) bool {
	if confRegion == "" {
		return true
	}
	if confRegion == "US" {
		confRegion = globalMinioDefaultRegion
	}
	// Some older s3 clients set region as "US" instead of
	// globalMinioDefaultRegion, handle it.
	if reqRegion == "US" {
		reqRegion = globalMinioDefaultRegion
	}
	return reqRegion == confRegion
}

// check if the access key is valid and recognized, additionally
// also returns if the access key is owner/admin.
func checkKeyValid(r *http.Request, accessKey string) (auth.Credentials, bool, APIErrorCode) {
	if !globalIAMSys.Initialized() && !globalIsGateway {
		// Check if server has initialized, then only proceed
		// to check for IAM users otherwise its okay for clients
		// to retry with 503 errors when server is coming up.
		return auth.Credentials{}, false, ErrServerNotInitialized
	}
	var owner = true
	var cred = globalActiveCred
	if cred.AccessKey != accessKey {
		// Check if the access key is part of users credentials.
		ucred, ok := globalIAMSys.GetUser(accessKey)
		if !ok {
			return cred, false, ErrInvalidAccessKeyID
		}
		claims, s3Err := checkClaimsFromToken(r, cred)
		if s3Err != ErrNone {
			return cred, false, s3Err
		}
		ucred.Claims = claims
		// Now check if we have a sessionPolicy.
		if _, ok = claims[iampolicy.SessionPolicyName]; ok {
			owner = false
		} else {
			owner = cred.AccessKey == ucred.ParentUser
		}
		cred = ucred
	}
	return cred, owner, ErrNone
}

// sumHMAC calculate hmac between two input byte array.
func sumHMAC(key []byte, data []byte) []byte {
	hash := hmac.New(sha256.New, key)
	hash.Write(data)
	return hash.Sum(nil)
}

// extractSignedHeaders extract signed headers from Authorization header
func extractSignedHeaders(signedHeaders []string, r *http.Request) (http.Header, APIErrorCode) {
	reqHeaders := r.Header
	reqQueries := r.Form
	// find whether "host" is part of list of signed headers.
	// if not return ErrUnsignedHeaders. "host" is mandatory.
	if !contains(signedHeaders, "host") {
		return nil, ErrUnsignedHeaders
	}
	extractedSignedHeaders := make(http.Header)
	for _, header := range signedHeaders {
		// `host` will not be found in the headers, can be found in r.Host.
		// but its alway necessary that the list of signed headers containing host in it.
		val, ok := reqHeaders[http.CanonicalHeaderKey(header)]
		if !ok {
			// try to set headers from Query String
			val, ok = reqQueries[header]
		}
		if ok {
			extractedSignedHeaders[http.CanonicalHeaderKey(header)] = val
			continue
		}
		switch header {
		case "expect":
			// Golang http server strips off 'Expect' header, if the
			// client sent this as part of signed headers we need to
			// handle otherwise we would see a signature mismatch.
			// `aws-cli` sets this as part of signed headers.
			//
			// According to
			// http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.20
			// Expect header is always of form:
			//
			//   Expect       =  "Expect" ":" 1#expectation
			//   expectation  =  "100-continue" | expectation-extension
			//
			// So it safe to assume that '100-continue' is what would
			// be sent, for the time being keep this work around.
			// Adding a *TODO* to remove this later when Golang server
			// doesn't filter out the 'Expect' header.
			extractedSignedHeaders.Set(header, "100-continue")
		case "host":
			// Go http server removes "host" from Request.Header
			extractedSignedHeaders.Set(header, r.Host)
		case "transfer-encoding":
			// Go http server removes "host" from Request.Header
			extractedSignedHeaders[http.CanonicalHeaderKey(header)] = r.TransferEncoding
		case "content-length":
			// Signature-V4 spec excludes Content-Length from signed headers list for signature calculation.
			// But some clients deviate from this rule. Hence we consider Content-Length for signature
			// calculation to be compatible with such clients.
			extractedSignedHeaders.Set(header, strconv.FormatInt(r.ContentLength, 10))
		default:
			return nil, ErrUnsignedHeaders
		}
	}
	return extractedSignedHeaders, ErrNone
}

// Trim leading and trailing spaces and replace sequential spaces with one space, following Trimall()
// in http://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html
func signV4TrimAll(input string) string {
	// Compress adjacent spaces (a space is determined by
	// unicode.IsSpace() internally here) to one space and return
	return strings.Join(strings.Fields(input), " ")
}
<commit_msg>pass the current credentials for claims<commit_after>// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	"github.com/minio/minio/internal/auth"
	xhttp "github.com/minio/minio/internal/http"
	"github.com/minio/minio/internal/logger"
	iampolicy "github.com/minio/pkg/iam/policy"
)

// http Header "x-amz-content-sha256" == "UNSIGNED-PAYLOAD" indicates that the
// client did not calculate sha256 of the payload.
const unsignedPayload = "UNSIGNED-PAYLOAD"

// skipContentSha256Cksum returns true if caller needs to skip
// payload checksum, false if not.
func skipContentSha256Cksum(r *http.Request) bool {
	var (
		v  []string
		ok bool
	)

	if isRequestPresignedSignatureV4(r) {
		v, ok = r.Form[xhttp.AmzContentSha256]
		if !ok {
			v, ok = r.Header[xhttp.AmzContentSha256]
		}
	} else {
		v, ok = r.Header[xhttp.AmzContentSha256]
	}

	// If x-amz-content-sha256 is set and the value is not
	// 'UNSIGNED-PAYLOAD' we should validate the content sha256.
	return !(ok && v[0] != unsignedPayload)
}

// Returns SHA256 for calculating canonical-request.
func getContentSha256Cksum(r *http.Request, stype serviceType) string {
	if stype == serviceSTS {
		payload, err := ioutil.ReadAll(io.LimitReader(r.Body, stsRequestBodyLimit))
		if err != nil {
			logger.CriticalIf(GlobalContext, err)
		}
		sum256 := sha256.Sum256(payload)
		r.Body = ioutil.NopCloser(bytes.NewReader(payload))
		return hex.EncodeToString(sum256[:])
	}

	var (
		defaultSha256Cksum string
		v                  []string
		ok                 bool
	)

	// For a presigned request we look at the query param for sha256.
	if isRequestPresignedSignatureV4(r) {
		// X-Amz-Content-Sha256, if not set in presigned requests, checksum
		// will default to 'UNSIGNED-PAYLOAD'.
		defaultSha256Cksum = unsignedPayload
		v, ok = r.Form[xhttp.AmzContentSha256]
		if !ok {
			v, ok = r.Header[xhttp.AmzContentSha256]
		}
	} else {
		// X-Amz-Content-Sha256, if not set in signed requests, checksum
		// will default to sha256([]byte("")).
		defaultSha256Cksum = emptySHA256
		v, ok = r.Header[xhttp.AmzContentSha256]
	}

	// We found 'X-Amz-Content-Sha256' return the captured value.
	if ok {
		return v[0]
	}

	// We couldn't find 'X-Amz-Content-Sha256'.
	return defaultSha256Cksum
}

// isValidRegion - verify if incoming region value is valid with configured Region.
func isValidRegion(reqRegion string, confRegion string) bool {
	if confRegion == "" {
		return true
	}
	if confRegion == "US" {
		confRegion = globalMinioDefaultRegion
	}
	// Some older s3 clients set region as "US" instead of
	// globalMinioDefaultRegion, handle it.
	if reqRegion == "US" {
		reqRegion = globalMinioDefaultRegion
	}
	return reqRegion == confRegion
}

// check if the access key is valid and recognized, additionally
// also returns if the access key is owner/admin.
func checkKeyValid(r *http.Request, accessKey string) (auth.Credentials, bool, APIErrorCode) {
	if !globalIAMSys.Initialized() && !globalIsGateway {
		// Check if server has initialized, then only proceed
		// to check for IAM users otherwise its okay for clients
		// to retry with 503 errors when server is coming up.
		return auth.Credentials{}, false, ErrServerNotInitialized
	}
	var owner = true
	var cred = globalActiveCred
	if cred.AccessKey != accessKey {
		// Check if the access key is part of users credentials.
		ucred, ok := globalIAMSys.GetUser(accessKey)
		if !ok {
			return cred, false, ErrInvalidAccessKeyID
		}
		claims, s3Err := checkClaimsFromToken(r, ucred)
		if s3Err != ErrNone {
			return cred, false, s3Err
		}
		ucred.Claims = claims
		// Now check if we have a sessionPolicy.
		if _, ok = claims[iampolicy.SessionPolicyName]; ok {
			owner = false
		} else {
			owner = cred.AccessKey == ucred.ParentUser
		}
		cred = ucred
	}
	return cred, owner, ErrNone
}

// sumHMAC calculate hmac between two input byte array.
func sumHMAC(key []byte, data []byte) []byte {
	hash := hmac.New(sha256.New, key)
	hash.Write(data)
	return hash.Sum(nil)
}

// extractSignedHeaders extract signed headers from Authorization header
func extractSignedHeaders(signedHeaders []string, r *http.Request) (http.Header, APIErrorCode) {
	reqHeaders := r.Header
	reqQueries := r.Form
	// find whether "host" is part of list of signed headers.
	// if not return ErrUnsignedHeaders. "host" is mandatory.
	if !contains(signedHeaders, "host") {
		return nil, ErrUnsignedHeaders
	}
	extractedSignedHeaders := make(http.Header)
	for _, header := range signedHeaders {
		// `host` will not be found in the headers, can be found in r.Host.
		// but its alway necessary that the list of signed headers containing host in it.
		val, ok := reqHeaders[http.CanonicalHeaderKey(header)]
		if !ok {
			// try to set headers from Query String
			val, ok = reqQueries[header]
		}
		if ok {
			extractedSignedHeaders[http.CanonicalHeaderKey(header)] = val
			continue
		}
		switch header {
		case "expect":
			// Golang http server strips off 'Expect' header, if the
			// client sent this as part of signed headers we need to
			// handle otherwise we would see a signature mismatch.
			// `aws-cli` sets this as part of signed headers.
			//
			// According to
			// http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.20
			// Expect header is always of form:
			//
			//   Expect       =  "Expect" ":" 1#expectation
			//   expectation  =  "100-continue" | expectation-extension
			//
			// So it safe to assume that '100-continue' is what would
			// be sent, for the time being keep this work around.
			// Adding a *TODO* to remove this later when Golang server
			// doesn't filter out the 'Expect' header.
			extractedSignedHeaders.Set(header, "100-continue")
		case "host":
			// Go http server removes "host" from Request.Header
			extractedSignedHeaders.Set(header, r.Host)
		case "transfer-encoding":
			// Go http server removes "host" from Request.Header
			extractedSignedHeaders[http.CanonicalHeaderKey(header)] = r.TransferEncoding
		case "content-length":
			// Signature-V4 spec excludes Content-Length from signed headers list for signature calculation.
			// But some clients deviate from this rule. Hence we consider Content-Length for signature
			// calculation to be compatible with such clients.
			extractedSignedHeaders.Set(header, strconv.FormatInt(r.ContentLength, 10))
		default:
			return nil, ErrUnsignedHeaders
		}
	}
	return extractedSignedHeaders, ErrNone
}

// Trim leading and trailing spaces and replace sequential spaces with one space, following Trimall()
// in http://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html
func signV4TrimAll(input string) string {
	// Compress adjacent spaces (a space is determined by
	// unicode.IsSpace() internally here) to one space and return
	return strings.Join(strings.Fields(input), " ")
}
<|endoftext|>


<commit_before>// Copyright 2014 Chadev. 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 (
	"strconv"
	"math/rand"

	"github.com/danryan/hal"
)

var cageMeHandler = hal.Hear(listenName+` cageme`, func(res *hal.Response) error {
	root := "http://cageme.herokuapp.com"
	num := strconv.Itoa(int(rand.Float64()*79+1))

	return res.Send(root + "/specific/" + num + ".jpeg")
})
<commit_msg>Run go fmt.<commit_after>// Copyright 2014 Chadev. 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 (
	"math/rand"
	"strconv"

	"github.com/danryan/hal"
)

var cageMeHandler = hal.Hear(listenName+` cageme`, func(res *hal.Response) error {
	root := "http://cageme.herokuapp.com"
	num := strconv.Itoa(int(rand.Float64()*79 + 1))

	return res.Send(root + "/specific/" + num + ".jpeg")
})
<|endoftext|>


<commit_before>// Copyright 2017-2018 Vasiliy Vasilyuk. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package blockchain // import "github.com/xorcare/blockchain"

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

const (
	// Version api client.
	Version = "1.0"

	// UserAgent is the header string used to identify this package.
	UserAgent = "blockchain-api-go-client/" + Version + " (go; github; +https://git.io/fNhcB)"

	// BasePath the root address in the network.
	BasePath = "https://blockchain.info"

	// BasePathTor the root address in the tor network.
	BasePathTor = "https://blockchainbdgpzk.onion"
)

// Errors it is a set of errors returned when working with the package.
var (
	ErrAIW = errors.New("address is wrong")
	ErrBEW = errors.New("block height is wrong")
	ErrBHW = errors.New("block hash is wrong")
	ErrCGD = errors.New("cannot get data on url")
	ErrCRR = errors.New("could not read answer response")
	ErrIRS = errors.New("incorrect response status")
	ErrNAP = errors.New("no address(es) provided")
	ErrRPE = errors.New("response parsing error")
	ErrTHW = errors.New("transaction hash is wrong")
)

// MaxAddressesCount the maximum number of addresses that can be checked at a time using
// the address checking function or the balance.
var MaxAddressesCount = 100

// Options map contains the default settings for requests to the api.
var Options = map[string]string{"format": "json"}

func init() {
	MaxAddressesCount = len(addressesForTestings)
}

// Client specifies the mechanism by which individual API requests are made.
type Client struct {
	client *http.Client

	APIKey    string // API access key.
	BasePath  string // API endpoint base URL.
	UserAgent string // Optional additional User-Agent fragment.
}

func (c *Client) userAgent() string {
	c.UserAgent = strings.TrimSpace(c.UserAgent)
	if c.UserAgent == "" {
		return UserAgent
	}

	return UserAgent + " " + c.UserAgent
}

func (c *Client) setErrorOne(errorMain error) error {
	return c.setError(errorMain, nil, nil, nil)
}

func (c *Client) setErrorTwo(errorMain error, errorExec error) error {
	return c.setError(errorMain, errorExec, nil, nil)
}

func (c *Client) setErrorThree(errorMain error, errorExec error, response *http.Response) error {
	return c.setError(errorMain, errorExec, response, nil)
}

func (c *Client) setError(errorMain error, errorExec error, response *http.Response, address *string) error {
	if errorMain == nil {
		return nil
	}

	return &Error{
		ErrMain:  errorMain,
		ErrExec:  errorExec,
		Response: response,
		Address:  address,
	}
}

// Do to send an client request, which is then converted to the passed type.
func (c *Client) Do(path string, i interface{}, options map[string]string) error {
	options = ApproveOptions(options)
	options["format"] = "json"
	options["api_code"] = c.APIKey
	values := url.Values{}
	for k, v := range options {
		values.Set(k, v)
	}

	req, e := http.NewRequest("GET", c.BasePath+path+"?"+(values.Encode()), nil)
	if e != nil {
		return c.setErrorTwo(ErrCGD, e)
	}
	req.Header.Set("User-Agent", c.userAgent())

	resp, e := c.client.Do(req)
	if e != nil {
		return c.setErrorThree(ErrCGD, e, resp)
	}
	defer resp.Body.Close()

	bytes, e := ioutil.ReadAll(resp.Body)
	if e != nil {
		return c.setErrorThree(ErrCRR, e, resp)
	}

	if resp.Status[0] != '2' {
		return c.setErrorThree(ErrIRS, errors.New(string(bytes)), resp)
	}

	if e = json.Unmarshal(bytes, &i); e != nil {
		return c.setErrorThree(ErrRPE, e, resp)
	}

	return nil
}

// Error data structure describing the error.
type Error struct {
	// Address wrong address.
	Address *string
	// ErrMain error information from the standard package error set.
	ErrMain error
	// ErrExec information about the error that occurred during
	// the operation of the standard library or external packages.
	ErrExec error
	// Response http response.
	Response *http.Response
}

// Error compatibility with error interface.
func (e Error) Error() string {
	return e.ErrMain.Error()
}

// ApproveOptions automatic check and substitution of options.
func ApproveOptions(options map[string]string) map[string]string {
	if options == nil {
		return Options
	}
	return options
}

// New creates a new client instance the network internet.
func New() *Client {
	return &Client{client: &http.Client{}, BasePath: BasePath}
}

// NewTor creates a new client instance the tor network.
func NewTor() *Client {
	return &Client{client: &http.Client{}, BasePath: BasePathTor}
}

// SetClient http client setter.
func (c *Client) SetClient(client *http.Client) {
	if client == nil {
		panic("blockchain: impossible install the client as nil")
	}
	c.client = client
}
<commit_msg>feat: Add package name to errors<commit_after>// Copyright 2017-2018 Vasiliy Vasilyuk. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package blockchain // import "github.com/xorcare/blockchain"

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

const (
	// Version api client.
	Version = "1.0"

	//Package package name const
	Package = "blockchain"

	// UserAgent is the header string used to identify this package.
	UserAgent = Package + "-api-go-client/" + Version + " (go; github; +https://git.io/fNhcB)"

	// BasePath the root address in the network.
	BasePath = "https://blockchain.info"

	// BasePathTor the root address in the tor network.
	BasePathTor = "https://blockchainbdgpzk.onion"
)

// Errors it is a set of errors returned when working with the package.
var (
	ErrAIW = errors.New(Package + ": address is wrong")
	ErrBEW = errors.New(Package + ": block height is wrong")
	ErrBHW = errors.New(Package + ": block hash is wrong")
	ErrCGD = errors.New(Package + ": cannot get data on url")
	ErrCRR = errors.New(Package + ": could not read answer response")
	ErrIRS = errors.New(Package + ": incorrect response status")
	ErrNAP = errors.New(Package + ": no address(es) provided")
	ErrRPE = errors.New(Package + ": response parsing error")
	ErrTHW = errors.New(Package + ": transaction hash is wrong")
)

// MaxAddressesCount the maximum number of addresses that can be checked at a time using
// the address checking function or the balance.
var MaxAddressesCount = 100

// Options map contains the default settings for requests to the api.
var Options = map[string]string{"format": "json"}

func init() {
	MaxAddressesCount = len(addressesForTestings)
}

// Client specifies the mechanism by which individual API requests are made.
type Client struct {
	client *http.Client

	APIKey    string // API access key.
	BasePath  string // API endpoint base URL.
	UserAgent string // Optional additional User-Agent fragment.
}

func (c *Client) userAgent() string {
	c.UserAgent = strings.TrimSpace(c.UserAgent)
	if c.UserAgent == "" {
		return UserAgent
	}

	return UserAgent + " " + c.UserAgent
}

func (c *Client) setErrorOne(errorMain error) error {
	return c.setError(errorMain, nil, nil, nil)
}

func (c *Client) setErrorTwo(errorMain error, errorExec error) error {
	return c.setError(errorMain, errorExec, nil, nil)
}

func (c *Client) setErrorThree(errorMain error, errorExec error, response *http.Response) error {
	return c.setError(errorMain, errorExec, response, nil)
}

func (c *Client) setError(errorMain error, errorExec error, response *http.Response, address *string) error {
	if errorMain == nil {
		return nil
	}

	return &Error{
		ErrMain:  errorMain,
		ErrExec:  errorExec,
		Response: response,
		Address:  address,
	}
}

// Do to send an client request, which is then converted to the passed type.
func (c *Client) Do(path string, i interface{}, options map[string]string) error {
	options = ApproveOptions(options)
	options["format"] = "json"
	options["api_code"] = c.APIKey
	values := url.Values{}
	for k, v := range options {
		values.Set(k, v)
	}

	req, e := http.NewRequest("GET", c.BasePath+path+"?"+(values.Encode()), nil)
	if e != nil {
		return c.setErrorTwo(ErrCGD, e)
	}
	req.Header.Set("User-Agent", c.userAgent())

	resp, e := c.client.Do(req)
	if e != nil {
		return c.setErrorThree(ErrCGD, e, resp)
	}
	defer resp.Body.Close()

	bytes, e := ioutil.ReadAll(resp.Body)
	if e != nil {
		return c.setErrorThree(ErrCRR, e, resp)
	}

	if resp.Status[0] != '2' {
		return c.setErrorThree(ErrIRS, errors.New(string(bytes)), resp)
	}

	if e = json.Unmarshal(bytes, &i); e != nil {
		return c.setErrorThree(ErrRPE, e, resp)
	}

	return nil
}

// Error data structure describing the error.
type Error struct {
	// Address wrong address.
	Address *string
	// ErrMain error information from the standard package error set.
	ErrMain error
	// ErrExec information about the error that occurred during
	// the operation of the standard library or external packages.
	ErrExec error
	// Response http response.
	Response *http.Response
}

// Error compatibility with error interface.
func (e Error) Error() string {
	return e.ErrMain.Error()
}

// ApproveOptions automatic check and substitution of options.
func ApproveOptions(options map[string]string) map[string]string {
	if options == nil {
		return Options
	}
	return options
}

// New creates a new client instance the network internet.
func New() *Client {
	return &Client{client: &http.Client{}, BasePath: BasePath}
}

// NewTor creates a new client instance the tor network.
func NewTor() *Client {
	return &Client{client: &http.Client{}, BasePath: BasePathTor}
}

// SetClient http client setter.
func (c *Client) SetClient(client *http.Client) {
	if client == nil {
		panic(Package + ": impossible install the client as nil")
	}
	c.client = client
}
<|endoftext|>


<commit_before>package git

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/tyba/oss/sources/vcs/clients/git/pktline"

	"github.com/sourcegraph/go-vcsurl"
)

type Client struct {
	url    string
	client *http.Client
}

func NewClient(url string) *Client {
	vcs, _ := vcsurl.Parse(url)
	return &Client{url: vcs.Link(), client: &http.Client{}}
}

func (c *Client) GetLastCommit() (string, error) {
	req, _ := c.buildRequest(
		"GET",
		fmt.Sprintf("%s/info/refs?service=git-upload-pack", c.url),
		nil,
	)

	res, err := c.client.Do(req)
	if err != nil {
		return "", err
	}

	if res.StatusCode >= 400 {
		return "", &NotFoundError{c.url}
	}

	defer res.Body.Close()
	d := pktline.NewDecoder(res.Body)

	content, err := d.ReadAll()
	if err != nil {
		return "", err
	}

	var head string
	for _, line := range content {
		if line[0] == '#' {
			continue
		}

		if head == "" {
			head = c.getHEADFromLine(line)
		} else {
			commit, branch := c.getCommitAndBranch(line)
			if branch == head {
				return commit, nil
			}
		}
	}

	return "", nil
}

func (c *Client) getHEADFromLine(line string) string {
	args, _ := url.ParseQuery(strings.Replace(line, " ", "&", -1))

	link, ok := args["symref"]
	if !ok {
		return ""
	}

	parts := strings.Split(link[0], ":")
	if len(parts) != 2 || parts[0] != "HEAD" {
		return ""
	}

	return parts[1]
}

func (c *Client) getCommitAndBranch(line string) (string, string) {
	parts := strings.Split(strings.Trim(line, " \n"), " ")
	if len(parts) != 2 {
		return "", ""
	}

	return parts[0], parts[1]
}

func (c *Client) GetPackFile(want string) (io.ReadCloser, error) {
	e := pktline.NewEncoder()
	e.AddLine(fmt.Sprintf("want %s", want))
	e.AddFlush()
	e.AddLine("done")

	req, err := c.buildRequest(
		"POST",
		fmt.Sprintf("%s/git-upload-pack", c.url),
		e.GetReader(),
	)
	if err != nil {
		return nil, err
	}

	res, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	h := make([]byte, 8)
	if _, err := res.Body.Read(h); err != nil {
		return nil, err
	}

	return res.Body, nil
}

func (c *Client) buildRequest(method, url string, content *strings.Reader) (*http.Request, error) {
	var req *http.Request
	var err error
	if content == nil {
		req, err = http.NewRequest(method, url, nil)
	} else {
		req, err = http.NewRequest(method, url, content)
	}

	if err != nil {
		return nil, err
	}

	req.Header.Add("User-Agent", "git/1.0")
	req.Header.Add("Host", "github.com")

	if content == nil {
		req.Header.Add("Accept", "*/*")
	} else {
		req.Header.Add("Accept", "application/x-git-upload-pack-result")
		req.Header.Add("Content-Type", "application/x-git-upload-pack-request")
		req.Header.Add("Content-Length", string(content.Len()))
	}

	return req, nil
}

type NotFoundError struct {
	url string
}

func (e NotFoundError) Error() string {
	return e.url
}
<commit_msg>Update imports<commit_after>package git

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/tyba/srcd-crawler/clients/git/pktline"

	"github.com/sourcegraph/go-vcsurl"
)

type Client struct {
	url    string
	client *http.Client
}

func NewClient(url string) *Client {
	vcs, _ := vcsurl.Parse(url)
	return &Client{url: vcs.Link(), client: &http.Client{}}
}

func (c *Client) GetLastCommit() (string, error) {
	req, _ := c.buildRequest(
		"GET",
		fmt.Sprintf("%s/info/refs?service=git-upload-pack", c.url),
		nil,
	)

	res, err := c.client.Do(req)
	if err != nil {
		return "", err
	}

	if res.StatusCode >= 400 {
		return "", &NotFoundError{c.url}
	}

	defer res.Body.Close()
	d := pktline.NewDecoder(res.Body)

	content, err := d.ReadAll()
	if err != nil {
		return "", err
	}

	var head string
	for _, line := range content {
		if line[0] == '#' {
			continue
		}

		if head == "" {
			head = c.getHEADFromLine(line)
		} else {
			commit, branch := c.getCommitAndBranch(line)
			if branch == head {
				return commit, nil
			}
		}
	}

	return "", nil
}

func (c *Client) getHEADFromLine(line string) string {
	args, _ := url.ParseQuery(strings.Replace(line, " ", "&", -1))

	link, ok := args["symref"]
	if !ok {
		return ""
	}

	parts := strings.Split(link[0], ":")
	if len(parts) != 2 || parts[0] != "HEAD" {
		return ""
	}

	return parts[1]
}

func (c *Client) getCommitAndBranch(line string) (string, string) {
	parts := strings.Split(strings.Trim(line, " \n"), " ")
	if len(parts) != 2 {
		return "", ""
	}

	return parts[0], parts[1]
}

func (c *Client) GetPackFile(want string) (io.ReadCloser, error) {
	e := pktline.NewEncoder()
	e.AddLine(fmt.Sprintf("want %s", want))
	e.AddFlush()
	e.AddLine("done")

	req, err := c.buildRequest(
		"POST",
		fmt.Sprintf("%s/git-upload-pack", c.url),
		e.GetReader(),
	)
	if err != nil {
		return nil, err
	}

	res, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	h := make([]byte, 8)
	if _, err := res.Body.Read(h); err != nil {
		return nil, err
	}

	return res.Body, nil
}

func (c *Client) buildRequest(method, url string, content *strings.Reader) (*http.Request, error) {
	var req *http.Request
	var err error
	if content == nil {
		req, err = http.NewRequest(method, url, nil)
	} else {
		req, err = http.NewRequest(method, url, content)
	}

	if err != nil {
		return nil, err
	}

	req.Header.Add("User-Agent", "git/1.0")
	req.Header.Add("Host", "github.com")

	if content == nil {
		req.Header.Add("Accept", "*/*")
	} else {
		req.Header.Add("Accept", "application/x-git-upload-pack-result")
		req.Header.Add("Content-Type", "application/x-git-upload-pack-request")
		req.Header.Add("Content-Length", string(content.Len()))
	}

	return req, nil
}

type NotFoundError struct {
	url string
}

func (e NotFoundError) Error() string {
	return e.url
}
<|endoftext|>


<commit_before>// Package caa provides access to the Cover Art Archive (http://coverartarchive.org)
package caa

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"path/filepath"
	"strconv"

	"code.google.com/p/go-uuid/uuid"
)

const baseurl = "http://coverartarchive.org"

// CAAClient manages the communication with the Cover Art Archive.
type CAAClient struct {
	useragent string
	client    http.Client
	BaseURL   string
}

// NewCAAClient returns a new CAAClient that uses the User-Agent useragent
func NewCAAClient(useragent string) (c *CAAClient) {
	c = &CAAClient{useragent: useragent, client: http.Client{}, BaseURL: baseurl}
	return
}

func (c *CAAClient) buildURL(path string) (url *url.URL) {
	url, err := url.Parse(c.BaseURL)

	if err != nil {
		return
	}

	url.Path = path
	return
}

func (c *CAAClient) get(url *url.URL) (resp *http.Response, err error) {
	req, _ := http.NewRequest("GET", url.String(), nil)
	req.Header.Set("User-Agent", c.useragent)

	resp, err = c.client.Do(req)

	if resp.StatusCode != http.StatusOK {
		err = HTTPError{StatusCode: resp.StatusCode, URL: url}
		return
	}

	if err != nil {
		log.Fatalln(err)
		return nil, err
	}

	return
}

func (c *CAAClient) getAndJSON(url *url.URL) (info *CoverArtInfo, err error) {
	resp, err := c.get(url)

	defer resp.Body.Close()

	if err != nil {
		return
	}

	body, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return
	}

	err = json.Unmarshal(body, &info)

	return

}

func (c *CAAClient) getImage(entitytype string, mbid uuid.UUID, imageid string, size int) (image CoverArtImage, err error) {
	var extra string

	if size == ImageSizeSmall || size == 250 {
		extra = "-250"
	} else if size == ImageSizeLarge || size == 500 {
		extra = "-500"
	} else {
		extra = ""
	}

	url := c.buildURL(entitytype + "/" + mbid.String() + "/" + imageid + extra)
	resp, err := c.get(url)

	defer resp.Body.Close()

	if err != nil {
		return
	}

	if resp.StatusCode != http.StatusOK {
		err = HTTPError{StatusCode: resp.StatusCode, URL: url}
		return
	}

	data, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return
	}

	image.Data = data

	ext := filepath.Ext(resp.Request.URL.String())
	mimetype := mime.TypeByExtension(ext)

	image.Mimetype = mimetype

	return

}

// GetReleaseInfo retrieves information about the images in the Cover Art Archive for the release with the MBID mbid
func (c *CAAClient) GetReleaseInfo(mbid uuid.UUID) (info *CoverArtInfo, err error) {
	url := c.buildURL("release/" + mbid.String())
	info, err = c.getAndJSON(url)
	return
}

// GetReleaseFront retrieves the front image of the release with the MBID mbid in the specified size
func (c *CAAClient) GetReleaseFront(mbid uuid.UUID, size int) (image CoverArtImage, err error) {
	image, err = c.getImage("release", mbid, "front", size)
	return
}

// GetReleaseBack retrieves the back image of the release with the MBID mbid in the specified size
func (c *CAAClient) GetReleaseBack(mbid uuid.UUID, size int) (image CoverArtImage, err error) {
	image, err = c.getImage("release", mbid, "back", size)
	return
}

// GetReleaseImage retrieves the image with the id imageid of the release with the MBID mbid in the specified size
func (c *CAAClient) GetReleaseImage(mbid uuid.UUID, imageid int, size int) (image CoverArtImage, err error) {
	id := strconv.Itoa(imageid)
	image, err = c.getImage("release", mbid, id, size)
	return
}

// GetReleaseGroupInfo retrieves information about the images in the Cover Art Archive for the release group with the MBID mbid
func (c *CAAClient) GetReleaseGroupInfo(mbid uuid.UUID) (info *CoverArtInfo, err error) {
	url := c.buildURL("release-group/" + mbid.String())
	info, err = c.getAndJSON(url)
	return
}

// GetReleaseGroupFront retrieves the front image of the release group with the MBID mbid in the specified size
func (c *CAAClient) GetReleaseGroupFront(mbid uuid.UUID, size int) (image CoverArtImage, err error) {
	if size != ImageSizeOriginal {
		err = InvalidImageSizeError{EntityType: "release-group", Size: size}
		return
	}
	image, err = c.getImage("release-group", mbid, "front", size)
	return
}
<commit_msg>Use HTTPS by default<commit_after>// Package caa provides access to the Cover Art Archive (http://coverartarchive.org)
package caa

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"path/filepath"
	"strconv"

	"code.google.com/p/go-uuid/uuid"
)

const baseurl = "https://coverartarchive.org"

// CAAClient manages the communication with the Cover Art Archive.
type CAAClient struct {
	useragent string
	client    http.Client
	BaseURL   string
}

// NewCAAClient returns a new CAAClient that uses the User-Agent useragent
func NewCAAClient(useragent string) (c *CAAClient) {
	c = &CAAClient{useragent: useragent, client: http.Client{}, BaseURL: baseurl}
	return
}

func (c *CAAClient) buildURL(path string) (url *url.URL) {
	url, err := url.Parse(c.BaseURL)

	if err != nil {
		return
	}

	url.Path = path
	return
}

func (c *CAAClient) get(url *url.URL) (resp *http.Response, err error) {
	req, _ := http.NewRequest("GET", url.String(), nil)
	req.Header.Set("User-Agent", c.useragent)

	resp, err = c.client.Do(req)

	if resp.StatusCode != http.StatusOK {
		err = HTTPError{StatusCode: resp.StatusCode, URL: url}
		return
	}

	if err != nil {
		log.Fatalln(err)
		return nil, err
	}

	return
}

func (c *CAAClient) getAndJSON(url *url.URL) (info *CoverArtInfo, err error) {
	resp, err := c.get(url)

	defer resp.Body.Close()

	if err != nil {
		return
	}

	body, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return
	}

	err = json.Unmarshal(body, &info)

	return

}

func (c *CAAClient) getImage(entitytype string, mbid uuid.UUID, imageid string, size int) (image CoverArtImage, err error) {
	var extra string

	if size == ImageSizeSmall || size == 250 {
		extra = "-250"
	} else if size == ImageSizeLarge || size == 500 {
		extra = "-500"
	} else {
		extra = ""
	}

	url := c.buildURL(entitytype + "/" + mbid.String() + "/" + imageid + extra)
	resp, err := c.get(url)

	defer resp.Body.Close()

	if err != nil {
		return
	}

	if resp.StatusCode != http.StatusOK {
		err = HTTPError{StatusCode: resp.StatusCode, URL: url}
		return
	}

	data, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return
	}

	image.Data = data

	ext := filepath.Ext(resp.Request.URL.String())
	mimetype := mime.TypeByExtension(ext)

	image.Mimetype = mimetype

	return

}

// GetReleaseInfo retrieves information about the images in the Cover Art Archive for the release with the MBID mbid
func (c *CAAClient) GetReleaseInfo(mbid uuid.UUID) (info *CoverArtInfo, err error) {
	url := c.buildURL("release/" + mbid.String())
	info, err = c.getAndJSON(url)
	return
}

// GetReleaseFront retrieves the front image of the release with the MBID mbid in the specified size
func (c *CAAClient) GetReleaseFront(mbid uuid.UUID, size int) (image CoverArtImage, err error) {
	image, err = c.getImage("release", mbid, "front", size)
	return
}

// GetReleaseBack retrieves the back image of the release with the MBID mbid in the specified size
func (c *CAAClient) GetReleaseBack(mbid uuid.UUID, size int) (image CoverArtImage, err error) {
	image, err = c.getImage("release", mbid, "back", size)
	return
}

// GetReleaseImage retrieves the image with the id imageid of the release with the MBID mbid in the specified size
func (c *CAAClient) GetReleaseImage(mbid uuid.UUID, imageid int, size int) (image CoverArtImage, err error) {
	id := strconv.Itoa(imageid)
	image, err = c.getImage("release", mbid, id, size)
	return
}

// GetReleaseGroupInfo retrieves information about the images in the Cover Art Archive for the release group with the MBID mbid
func (c *CAAClient) GetReleaseGroupInfo(mbid uuid.UUID) (info *CoverArtInfo, err error) {
	url := c.buildURL("release-group/" + mbid.String())
	info, err = c.getAndJSON(url)
	return
}

// GetReleaseGroupFront retrieves the front image of the release group with the MBID mbid in the specified size
func (c *CAAClient) GetReleaseGroupFront(mbid uuid.UUID, size int) (image CoverArtImage, err error) {
	if size != ImageSizeOriginal {
		err = InvalidImageSizeError{EntityType: "release-group", Size: size}
		return
	}
	image, err = c.getImage("release-group", mbid, "front", size)
	return
}
<|endoftext|>


<commit_before>package wsevent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn    *ws.Conn
	send    chan []byte
	recv    chan []byte
	request *http.Request

	close chan struct{}
}

type request struct {
	Id   string
	Data json.RawMessage
}

type reply struct {
	Id   string      `json:"id"`
	Data interface{} `json:"data"`
}

var (
	reqPool   = &sync.Pool{New: func() interface{} { return request{} }}
	replyPool = &sync.Pool{New: func() interface{} { return reply{} }}
)

func genID() string {
	bytes := make([]byte, 32)
	rand.Read(bytes)

	return base64.URLEncoding.EncodeToString(bytes)
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClientWithID(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request, id string) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:      id,
		conn:    conn,
		send:    make(chan []byte),
		request: r,

		close: make(chan struct{}, 3),
	}
	s.newClient <- client

	return client, nil
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	return s.NewClientWithID(upgrader, w, r, genID())
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) {
	c.send <- []byte(data)
}

type emitJS struct {
	Id   int         `json:"id"`
	Data interface{} `json:"data"`
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	js := emitJS{}
	js.Id = -1
	js.Data = v

	bytes, err := json.Marshal(js)
	if err != nil {
		return err
	}

	c.send <- bytes
	return nil
}

func (c *Client) Close() {
	for i := 0; i < 3; i++ {
		c.close <- struct{}{}
	}

	c.conn.Close()
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		s.roomsLock.Lock()
		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
		s.roomsLock.Unlock()
	}
	s.joinedRoomsLock.RUnlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

func (c *Client) listener(s *Server) {
	tick := time.NewTicker(time.Millisecond * 10)
	go func() {
		for {
			select {
			case data := <-c.send:
				c.conn.WriteMessage(ws.TextMessage, data)

			case <-c.close:
				return
			}
		}
	}()
	go func() {
		for {
			select {
			case <-tick.C:
				mtype, data, err := c.conn.ReadMessage()
				if err != nil {
					c.cleanup(s)
					c.recv <- []byte{}
					c.close <- struct{}{}
					tick.Stop()
					return
				}
				if mtype != ws.TextMessage {
					continue
				}

				c.recv <- data

			case <-c.close:
				return
			}
		}
	}()

	for {
		select {
		case data := <-c.recv:
			if len(data) == 0 {
				return
			}

			req := reqPool.Get().(request)

			if err := json.Unmarshal(data, &req); err != nil {
				continue
			}

			callName := s.Extractor(req.Data)

			s.handlersLock.RLock()
			f, ok := s.handlers[callName]
			s.handlersLock.RUnlock()

			if !ok {
				if s.DefaultHandler != nil {
					f = s.DefaultHandler
					goto call
				}
				continue
			}
		call:
			go func() {
				rtrn := f(s, c, req.Data)
				reply := replyPool.Get().(reply)
				reply.Id = req.Id
				reply.Data = rtrn

				bytes, _ := json.Marshal(reply)
				c.send <- bytes

				reqPool.Put(req)
				replyPool.Put(reply)
			}()

		case <-c.close:
			return
		}
	}
}
<commit_msg>Make recv channel.<commit_after>package wsevent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn    *ws.Conn
	send    chan []byte
	recv    chan []byte
	request *http.Request

	close chan struct{}
}

type request struct {
	Id   string
	Data json.RawMessage
}

type reply struct {
	Id   string      `json:"id"`
	Data interface{} `json:"data"`
}

var (
	reqPool   = &sync.Pool{New: func() interface{} { return request{} }}
	replyPool = &sync.Pool{New: func() interface{} { return reply{} }}
)

func genID() string {
	bytes := make([]byte, 32)
	rand.Read(bytes)

	return base64.URLEncoding.EncodeToString(bytes)
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClientWithID(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request, id string) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:      id,
		conn:    conn,
		send:    make(chan []byte),
		recv:    make(chan []byte),
		request: r,

		close: make(chan struct{}, 3),
	}
	s.newClient <- client

	return client, nil
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	return s.NewClientWithID(upgrader, w, r, genID())
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) {
	c.send <- []byte(data)
}

type emitJS struct {
	Id   int         `json:"id"`
	Data interface{} `json:"data"`
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	js := emitJS{}
	js.Id = -1
	js.Data = v

	bytes, err := json.Marshal(js)
	if err != nil {
		return err
	}

	c.send <- bytes
	return nil
}

func (c *Client) Close() {
	for i := 0; i < 3; i++ {
		c.close <- struct{}{}
	}

	c.conn.Close()
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		s.roomsLock.Lock()
		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
		s.roomsLock.Unlock()
	}
	s.joinedRoomsLock.RUnlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

func (c *Client) listener(s *Server) {
	tick := time.NewTicker(time.Millisecond * 10)
	go func() {
		for {
			select {
			case data := <-c.send:
				c.conn.WriteMessage(ws.TextMessage, data)

			case <-c.close:
				return
			}
		}
	}()
	go func() {
		for {
			select {
			case <-tick.C:
				mtype, data, err := c.conn.ReadMessage()
				if err != nil {
					c.cleanup(s)
					c.recv <- []byte{}
					c.close <- struct{}{}
					tick.Stop()
					return
				}
				if mtype != ws.TextMessage {
					continue
				}

				c.recv <- data

			case <-c.close:
				return
			}
		}
	}()

	for {
		select {
		case data := <-c.recv:
			if len(data) == 0 {
				return
			}

			req := reqPool.Get().(request)

			if err := json.Unmarshal(data, &req); err != nil {
				continue
			}

			callName := s.Extractor(req.Data)

			s.handlersLock.RLock()
			f, ok := s.handlers[callName]
			s.handlersLock.RUnlock()

			if !ok {
				if s.DefaultHandler != nil {
					f = s.DefaultHandler
					goto call
				}
				continue
			}
		call:
			go func() {
				rtrn := f(s, c, req.Data)
				reply := replyPool.Get().(reply)
				reply.Id = req.Id
				reply.Data = rtrn

				bytes, _ := json.Marshal(reply)
				c.send <- bytes

				reqPool.Put(req)
				replyPool.Put(reply)
			}()

		case <-c.close:
			return
		}
	}
}
<|endoftext|>


<commit_before>/*
A smart client for go.

Usage:

 client, err := couchbase.Connect("http://myserver:8091/")
 handleError(err)
 pool, err := client.GetPool("default")
 handleError(err)
 bucket, err := pool.getBucket("MyAwesomeBucket")
 handleError(err)
 ...

or a shortcut for the bucket directly

 bucket, err := couchbase.GetBucket("http://myserver:8091/", "default", "default")
*/
package couchbase

import (
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

type connectionPool struct {
	host        string
	name        string
	connections []*memcached.Client
	mutex       sync.Mutex
}

func (cp *connectionPool) Close() error {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()
	for _, c := range cp.connections {
		c.Close()
	}
	cp.connections = []*memcached.Client{}
	return nil
}

func (cp *connectionPool) Get() (*memcached.Client, error) {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if len(cp.connections) == 0 {
		conn, err := memcached.Connect("tcp", cp.host)
		if err != nil {
			return nil, err
		}
		if cp.name != "default" {
			conn.Auth(cp.name, "")
		}

		cp.connections = append(cp.connections, conn)
	}

	rv := cp.connections[0]
	cp.connections = cp.connections[1:]

	return rv, nil
}

func (cp *connectionPool) Return(c *memcached.Client) {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if c != nil {
		if c.IsHealthy() {
			cp.connections = append(cp.connections, c)
		} else {
			c.Close()
		}
	}
}

// Execute a function on a memcached connection to the node owning key "k"
//
// Note that this automatically handles transient errors by replaying
// your function on a "not-my-vbucket" error, so don't assume
// your command will only be executed only once.
func (b *Bucket) Do(k string, f func(mc *memcached.Client, vb uint16) error) error {
	vb := b.VBHash(k)
	for {
		masterId := b.VBucketServerMap.VBucketMap[vb][0]
		conn, err := b.connections[masterId].Get()
		defer b.connections[masterId].Return(conn)
		if err != nil {
			return err
		}

		err = f(conn, uint16(vb))
		switch err.(type) {
		default:
			return err
		case gomemcached.MCResponse:
			st := err.(gomemcached.MCResponse).Status
			atomic.AddUint64(&b.pool.client.Statuses[st], 1)
			if st == gomemcached.NOT_MY_VBUCKET {
				b.refresh()
			} else {
				return err
			}
		}
	}
	panic("Unreachable.")
}

type gathered_stats struct {
	sn   string
	vals map[string]string
}

func getStatsParallel(b *Bucket, offset int, which string, ch chan<- gathered_stats) {
	sn := b.VBucketServerMap.ServerList[offset]

	results := map[string]string{}
	conn, err := b.connections[offset].Get()
	defer b.connections[offset].Return(conn)
	if err != nil {
		ch <- gathered_stats{sn, results}
	} else {
		st, err := conn.StatsMap(which)
		if err == nil {
			ch <- gathered_stats{sn, st}
		} else {
			ch <- gathered_stats{sn, results}
		}
	}
}

// Get a set of stats from all servers.
//
// Returns a map of server ID -> map of stat key to map value.
func (b *Bucket) GetStats(which string) map[string]map[string]string {
	rv := map[string]map[string]string{}

	if b.VBucketServerMap.ServerList == nil {
		return rv
	}
	// Go grab all the things at once.
	todo := len(b.VBucketServerMap.ServerList)
	ch := make(chan gathered_stats, todo)

	for offset, _ := range b.VBucketServerMap.ServerList {
		go getStatsParallel(b, offset, which, ch)
	}

	// Gather the results
	for i := 0; i < len(b.VBucketServerMap.ServerList); i++ {
		g := <-ch
		if len(g.vals) > 0 {
			rv[g.sn] = g.vals
		}
	}

	return rv
}

func (b *Bucket) doBulkGet(vb uint16, keys []string,
	ch chan<- map[string]*gomemcached.MCResponse) {

	masterId := b.VBucketServerMap.VBucketMap[vb][0]
	conn, err := b.connections[masterId].Get()
	if err != nil {
		ch <- map[string]*gomemcached.MCResponse{}
	}
	defer b.connections[masterId].Return(conn)

	m, err := conn.GetBulk(vb, keys)
	switch err.(type) {
	default:
		ch <- m
	case *gomemcached.MCResponse:
		fmt.Printf("Got a memcached error")
		st := err.(gomemcached.MCResponse).Status
		atomic.AddUint64(&b.pool.client.Statuses[st], 1)
		if st == gomemcached.NOT_MY_VBUCKET {
			b.refresh()
		}
		ch <- map[string]*gomemcached.MCResponse{}
	}
}

func (b *Bucket) processBulkGet(kdm map[uint16][]string,
	ch chan map[string]*gomemcached.MCResponse) {

	wch := make(chan uint16)

	worker := func() {
		for k := range wch {
			b.doBulkGet(k, kdm[k], ch)
		}
	}

	for i := 0; i < 4; i++ {
		go worker()
	}

	for k := range kdm {
		wch <- k
	}

}
func (b *Bucket) GetBulk(keys []string) map[string]*gomemcached.MCResponse {
	// Organize by vbucket
	kdm := map[uint16][]string{}
	for _, k := range keys {
		vb := uint16(b.VBHash(k))
		a, ok := kdm[vb]
		if !ok {
			a = []string{}
		}
		kdm[vb] = append(a, k)
	}

	ch := make(chan map[string]*gomemcached.MCResponse)
	defer close(ch)

	go b.processBulkGet(kdm, ch)

	rv := map[string]*gomemcached.MCResponse{}
	for _ = range kdm {
		m := <-ch
		for k, v := range m {
			rv[k] = v
		}
	}

	return rv
}

// Set a value in this bucket.
// The value will be serialized into a JSON document.
func (b *Bucket) Set(k string, v interface{}) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		data, err := json.Marshal(v)
		if err != nil {
			return err
		}
		res, err := mc.Set(vb, k, 0, 0, data)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		return nil
	})
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Gets(k string, rv interface{}, cas *uint64) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Get(vb, k)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		if cas != nil {
			*cas = res.Cas
		}
		return json.Unmarshal(res.Body, rv)
	})
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Get(k string, rv interface{}) error {
	return b.Gets(k, rv, nil)
}

// Delete a key from this bucket.
func (b *Bucket) Delete(k string) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Del(vb, k)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		return nil
	})
}

// Increment a key
func (b *Bucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	var rv uint64
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Incr(vb, k, amt, def, exp)
		if err != nil {
			return err
		}
		rv = res
		return nil
	})
	return rv, err
}

type ViewRow struct {
	ID    string
	Key   interface{}
	Value interface{}
	Doc   *interface{}
}

type ViewResult struct {
	TotalRows int `json:"total_rows"`
	Rows      []ViewRow
	Errors    []struct {
		From   string
		Reason string
	}
}

// Perform a view request that can map row values to a custom type.
//
// See the source to View for an example usage.
func (b *Bucket) ViewCustom(ddoc, name string, params map[string]interface{},
	vres interface{}) error {

	// Pick a random node to service our request.
	node := b.Nodes[rand.Intn(len(b.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		return err
	}

	values := url.Values{}
	for k, v := range params {
		switch t := v.(type) {
		case string:
			values[k] = []string{fmt.Sprintf(`"%s"`, url.QueryEscape(t))}
		case int:
			values[k] = []string{fmt.Sprintf(`%d`, t)}
		case bool:
			values[k] = []string{fmt.Sprintf(`%v`, t)}
		default:
			b, err := json.Marshal(v)
			if err != nil {
				panic(fmt.Sprintf("unsupported value-type %T in Query, json encoder said %v", t, err))
			}
			values[k] = []string{fmt.Sprintf(`%v`, string(b))}
		}
	}

	u.Path = fmt.Sprintf("/%s/_design/%s/_view/%s", b.Name, ddoc, name)
	u.RawQuery = values.Encode()

	res, err := http.Get(u.String())
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return errors.New(res.Status)
	}

	d := json.NewDecoder(res.Body)
	if err := d.Decode(vres); err != nil {
		return err
	}
	return nil
}

// Execute a view
func (b *Bucket) View(ddoc, name string, params map[string]interface{}) (ViewResult, error) {
	vres := ViewResult{}
	return vres, b.ViewCustom(ddoc, name, params, &vres)
}
<commit_msg>Take an expiration argument to set.<commit_after>/*
A smart client for go.

Usage:

 client, err := couchbase.Connect("http://myserver:8091/")
 handleError(err)
 pool, err := client.GetPool("default")
 handleError(err)
 bucket, err := pool.getBucket("MyAwesomeBucket")
 handleError(err)
 ...

or a shortcut for the bucket directly

 bucket, err := couchbase.GetBucket("http://myserver:8091/", "default", "default")
*/
package couchbase

import (
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

type connectionPool struct {
	host        string
	name        string
	connections []*memcached.Client
	mutex       sync.Mutex
}

func (cp *connectionPool) Close() error {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()
	for _, c := range cp.connections {
		c.Close()
	}
	cp.connections = []*memcached.Client{}
	return nil
}

func (cp *connectionPool) Get() (*memcached.Client, error) {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if len(cp.connections) == 0 {
		conn, err := memcached.Connect("tcp", cp.host)
		if err != nil {
			return nil, err
		}
		if cp.name != "default" {
			conn.Auth(cp.name, "")
		}

		cp.connections = append(cp.connections, conn)
	}

	rv := cp.connections[0]
	cp.connections = cp.connections[1:]

	return rv, nil
}

func (cp *connectionPool) Return(c *memcached.Client) {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if c != nil {
		if c.IsHealthy() {
			cp.connections = append(cp.connections, c)
		} else {
			c.Close()
		}
	}
}

// Execute a function on a memcached connection to the node owning key "k"
//
// Note that this automatically handles transient errors by replaying
// your function on a "not-my-vbucket" error, so don't assume
// your command will only be executed only once.
func (b *Bucket) Do(k string, f func(mc *memcached.Client, vb uint16) error) error {
	vb := b.VBHash(k)
	for {
		masterId := b.VBucketServerMap.VBucketMap[vb][0]
		conn, err := b.connections[masterId].Get()
		defer b.connections[masterId].Return(conn)
		if err != nil {
			return err
		}

		err = f(conn, uint16(vb))
		switch err.(type) {
		default:
			return err
		case gomemcached.MCResponse:
			st := err.(gomemcached.MCResponse).Status
			atomic.AddUint64(&b.pool.client.Statuses[st], 1)
			if st == gomemcached.NOT_MY_VBUCKET {
				b.refresh()
			} else {
				return err
			}
		}
	}
	panic("Unreachable.")
}

type gathered_stats struct {
	sn   string
	vals map[string]string
}

func getStatsParallel(b *Bucket, offset int, which string, ch chan<- gathered_stats) {
	sn := b.VBucketServerMap.ServerList[offset]

	results := map[string]string{}
	conn, err := b.connections[offset].Get()
	defer b.connections[offset].Return(conn)
	if err != nil {
		ch <- gathered_stats{sn, results}
	} else {
		st, err := conn.StatsMap(which)
		if err == nil {
			ch <- gathered_stats{sn, st}
		} else {
			ch <- gathered_stats{sn, results}
		}
	}
}

// Get a set of stats from all servers.
//
// Returns a map of server ID -> map of stat key to map value.
func (b *Bucket) GetStats(which string) map[string]map[string]string {
	rv := map[string]map[string]string{}

	if b.VBucketServerMap.ServerList == nil {
		return rv
	}
	// Go grab all the things at once.
	todo := len(b.VBucketServerMap.ServerList)
	ch := make(chan gathered_stats, todo)

	for offset, _ := range b.VBucketServerMap.ServerList {
		go getStatsParallel(b, offset, which, ch)
	}

	// Gather the results
	for i := 0; i < len(b.VBucketServerMap.ServerList); i++ {
		g := <-ch
		if len(g.vals) > 0 {
			rv[g.sn] = g.vals
		}
	}

	return rv
}

func (b *Bucket) doBulkGet(vb uint16, keys []string,
	ch chan<- map[string]*gomemcached.MCResponse) {

	masterId := b.VBucketServerMap.VBucketMap[vb][0]
	conn, err := b.connections[masterId].Get()
	if err != nil {
		ch <- map[string]*gomemcached.MCResponse{}
	}
	defer b.connections[masterId].Return(conn)

	m, err := conn.GetBulk(vb, keys)
	switch err.(type) {
	default:
		ch <- m
	case *gomemcached.MCResponse:
		fmt.Printf("Got a memcached error")
		st := err.(gomemcached.MCResponse).Status
		atomic.AddUint64(&b.pool.client.Statuses[st], 1)
		if st == gomemcached.NOT_MY_VBUCKET {
			b.refresh()
		}
		ch <- map[string]*gomemcached.MCResponse{}
	}
}

func (b *Bucket) processBulkGet(kdm map[uint16][]string,
	ch chan map[string]*gomemcached.MCResponse) {

	wch := make(chan uint16)

	worker := func() {
		for k := range wch {
			b.doBulkGet(k, kdm[k], ch)
		}
	}

	for i := 0; i < 4; i++ {
		go worker()
	}

	for k := range kdm {
		wch <- k
	}

}
func (b *Bucket) GetBulk(keys []string) map[string]*gomemcached.MCResponse {
	// Organize by vbucket
	kdm := map[uint16][]string{}
	for _, k := range keys {
		vb := uint16(b.VBHash(k))
		a, ok := kdm[vb]
		if !ok {
			a = []string{}
		}
		kdm[vb] = append(a, k)
	}

	ch := make(chan map[string]*gomemcached.MCResponse)
	defer close(ch)

	go b.processBulkGet(kdm, ch)

	rv := map[string]*gomemcached.MCResponse{}
	for _ = range kdm {
		m := <-ch
		for k, v := range m {
			rv[k] = v
		}
	}

	return rv
}

// Set a value in this bucket.
// The value will be serialized into a JSON document.
func (b *Bucket) Set(k string, exp int, v interface{}) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		data, err := json.Marshal(v)
		if err != nil {
			return err
		}
		res, err := mc.Set(vb, k, 0, exp, data)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		return nil
	})
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Gets(k string, rv interface{}, cas *uint64) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Get(vb, k)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		if cas != nil {
			*cas = res.Cas
		}
		return json.Unmarshal(res.Body, rv)
	})
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Get(k string, rv interface{}) error {
	return b.Gets(k, rv, nil)
}

// Delete a key from this bucket.
func (b *Bucket) Delete(k string) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Del(vb, k)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		return nil
	})
}

// Increment a key
func (b *Bucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	var rv uint64
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Incr(vb, k, amt, def, exp)
		if err != nil {
			return err
		}
		rv = res
		return nil
	})
	return rv, err
}

type ViewRow struct {
	ID    string
	Key   interface{}
	Value interface{}
	Doc   *interface{}
}

type ViewResult struct {
	TotalRows int `json:"total_rows"`
	Rows      []ViewRow
	Errors    []struct {
		From   string
		Reason string
	}
}

// Perform a view request that can map row values to a custom type.
//
// See the source to View for an example usage.
func (b *Bucket) ViewCustom(ddoc, name string, params map[string]interface{},
	vres interface{}) error {

	// Pick a random node to service our request.
	node := b.Nodes[rand.Intn(len(b.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		return err
	}

	values := url.Values{}
	for k, v := range params {
		switch t := v.(type) {
		case string:
			values[k] = []string{fmt.Sprintf(`"%s"`, url.QueryEscape(t))}
		case int:
			values[k] = []string{fmt.Sprintf(`%d`, t)}
		case bool:
			values[k] = []string{fmt.Sprintf(`%v`, t)}
		default:
			b, err := json.Marshal(v)
			if err != nil {
				panic(fmt.Sprintf("unsupported value-type %T in Query, json encoder said %v", t, err))
			}
			values[k] = []string{fmt.Sprintf(`%v`, string(b))}
		}
	}

	u.Path = fmt.Sprintf("/%s/_design/%s/_view/%s", b.Name, ddoc, name)
	u.RawQuery = values.Encode()

	res, err := http.Get(u.String())
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return errors.New(res.Status)
	}

	d := json.NewDecoder(res.Body)
	if err := d.Decode(vres); err != nil {
		return err
	}
	return nil
}

// Execute a view
func (b *Bucket) View(ddoc, name string, params map[string]interface{}) (ViewResult, error) {
	vres := ViewResult{}
	return vres, b.ViewCustom(ddoc, name, params, &vres)
}
<|endoftext|>