<commit_before>/*
   Copyright The containerd Authors.

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

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

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

package main

import (
	"context"
	"flag"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"

	"github.com/BurntSushi/toml"
	snapshotsapi "github.com/containerd/containerd/api/services/snapshots/v1"
	"github.com/containerd/containerd/contrib/snapshotservice"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/remotes/docker"
	"github.com/containerd/stargz-snapshotter/cmd/containerd-stargz-grpc/keychain"
	stargzfs "github.com/containerd/stargz-snapshotter/fs"
	"github.com/containerd/stargz-snapshotter/fs/source"
	snbase "github.com/containerd/stargz-snapshotter/snapshot"
	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/hashicorp/go-multierror"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"
)

const (
	defaultAddress  = "/run/containerd-stargz-grpc/containerd-stargz-grpc.sock"
	defaultLogLevel = logrus.InfoLevel
	defaultRootDir  = "/var/lib/containerd-stargz-grpc"
)

var (
	address    = flag.String("address", defaultAddress, "address for the snapshotter's GRPC server")
	configPath = flag.String("config", "", "path to the configuration file")
	logLevel   = flag.String("log-level", defaultLogLevel.String(), "set the logging level [trace, debug, info, warn, error, fatal, panic]")
	rootDir    = flag.String("root", defaultRootDir, "path to the root directory for this snapshotter")
)

func main() {
	flag.Parse()
	lvl, err := logrus.ParseLevel(*logLevel)
	if err != nil {
		log.L.WithError(err).Fatal("failed to prepare logger")
	}
	logrus.SetLevel(lvl)
	logrus.SetFormatter(&logrus.JSONFormatter{
		TimestampFormat: log.RFC3339NanoFixed,
	})

	var (
		ctx    = log.WithLogger(context.Background(), log.L)
		config Config
	)

	// Get configuration from specified file
	if *configPath != "" {
		if _, err := toml.DecodeFile(*configPath, &config); err != nil {
			log.G(ctx).WithError(err).Fatalf("failed to load config file %q", *configPath)
		}
	}

	// Prepare kubeconfig-based keychain if required
	kc := authn.DefaultKeychain
	if config.KubeconfigKeychainConfig.EnableKeychain {
		var opts []keychain.KubeconfigOption
		if kcp := config.KubeconfigKeychainConfig.KubeconfigPath; kcp != "" {
			opts = append(opts, keychain.WithKubeconfigPath(kcp))
		}
		kc = authn.NewMultiKeychain(kc, keychain.NewKubeconfigKeychain(ctx, opts...))
	}

	// Use RegistryHosts based on ResolverConfig and keychain
	hosts := hostsFromConfig(config.ResolverConfig, kc)

	// Configure filesystem and snapshotter
	fs, err := stargzfs.NewFilesystem(filepath.Join(*rootDir, "stargz"),
		config.Config,
		stargzfs.WithGetSources(sources(
			sourceFromCRILabels(hosts),      // provides source info based on CRI labels
			source.FromDefaultLabels(hosts), // provides source info based on default labels
		)),
	)
	if err != nil {
		log.G(ctx).WithError(err).Fatalf("failed to configure filesystem")
	}
	rs, err := snbase.NewSnapshotter(ctx, filepath.Join(*rootDir, "snapshotter"), fs, snbase.AsynchronousRemove)
	if err != nil {
		log.G(ctx).WithError(err).Fatalf("failed to configure snapshotter")
	}
	defer func() {
		log.G(ctx).Debug("Closing the snapshotter")
		rs.Close()
		log.G(ctx).Info("Exiting")
	}()

	// Create a gRPC server
	rpc := grpc.NewServer()

	// Convert the snapshotter to a gRPC service,
	service := snapshotservice.FromSnapshotter(rs)

	// Register the service with the gRPC server
	snapshotsapi.RegisterSnapshotsServer(rpc, service)

	// Prepare the directory for the socket
	if err := os.MkdirAll(filepath.Dir(*address), 0700); err != nil {
		log.G(ctx).WithError(err).Fatalf("failed to create directory %q", filepath.Dir(*address))
	}

	// Try to remove the socket file to avoid EADDRINUSE
	if err := os.RemoveAll(*address); err != nil {
		log.G(ctx).WithError(err).Fatalf("failed to remove %q", *address)
	}

	// Listen and serve
	l, err := net.Listen("unix", *address)
	if err != nil {
		log.G(ctx).WithError(err).Fatalf("error on listen socket %q", *address)
	}
	go func() {
		if err := rpc.Serve(l); err != nil {
			log.G(ctx).WithError(err).Fatalf("error on serving via socket %q", *address)
		}
	}()
	waitForSIGINT()
	log.G(ctx).Info("Got SIGINT")
}

func waitForSIGINT() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c
}

func hostsFromConfig(cfg ResolverConfig, keychain authn.Keychain) docker.RegistryHosts {
	return func(host string) (hosts []docker.RegistryHost, _ error) {
		for _, h := range append(cfg.Host[host].Mirrors, MirrorConfig{
			Host: host,
		}) {
			tr := &http.Client{Transport: http.DefaultTransport.(*http.Transport).Clone()}
			config := docker.RegistryHost{
				Client:       tr,
				Host:         h.Host,
				Scheme:       "https",
				Path:         "/v2",
				Capabilities: docker.HostCapabilityPull,
				Authorizer: docker.NewDockerAuthorizer(
					docker.WithAuthClient(tr),
					docker.WithAuthCreds(func(host string) (string, string, error) {
						if host == "registry-1.docker.io" {
							host = "index.docker.io"
						}
						reg, err := name.NewRegistry(host)
						if err != nil {
							return "", "", err
						}
						authn, err := keychain.Resolve(reg)
						if err != nil {
							return "", "", err
						}
						acfg, err := authn.Authorization()
						if err != nil {
							return "", "", err
						}
						if acfg.IdentityToken != "" {
							return "", acfg.IdentityToken, nil
						}
						return acfg.Username, acfg.Password, nil
					})),
			}
			if localhost, _ := docker.MatchLocalhost(config.Host); localhost || h.Insecure {
				config.Scheme = "http"
			}
			if config.Host == "docker.io" {
				config.Host = "registry-1.docker.io"
			}
			hosts = append(hosts, config)
		}
		return
	}
}

func sources(ps ...source.GetSources) source.GetSources {
	return func(labels map[string]string) (source []source.Source, allErr error) {
		for _, p := range ps {
			src, err := p(labels)
			if err == nil {
				return src, nil
			}
			allErr = multierror.Append(allErr, err)
		}
		return
	}
}
<commit_msg>Add defaultConfigPath<commit_after>/*
   Copyright The containerd Authors.

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

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

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

package main

import (
	"context"
	"flag"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"

	"github.com/BurntSushi/toml"
	snapshotsapi "github.com/containerd/containerd/api/services/snapshots/v1"
	"github.com/containerd/containerd/contrib/snapshotservice"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/remotes/docker"
	"github.com/containerd/stargz-snapshotter/cmd/containerd-stargz-grpc/keychain"
	stargzfs "github.com/containerd/stargz-snapshotter/fs"
	"github.com/containerd/stargz-snapshotter/fs/source"
	snbase "github.com/containerd/stargz-snapshotter/snapshot"
	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/hashicorp/go-multierror"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"
)

const (
	defaultAddress    = "/run/containerd-stargz-grpc/containerd-stargz-grpc.sock"
	defaultConfigPath = "/etc/containerd-stargz-grpc/config.toml"
	defaultLogLevel   = logrus.InfoLevel
	defaultRootDir    = "/var/lib/containerd-stargz-grpc"
)

var (
	address    = flag.String("address", defaultAddress, "address for the snapshotter's GRPC server")
	configPath = flag.String("config", defaultConfigPath, "path to the configuration file")
	logLevel   = flag.String("log-level", defaultLogLevel.String(), "set the logging level [trace, debug, info, warn, error, fatal, panic]")
	rootDir    = flag.String("root", defaultRootDir, "path to the root directory for this snapshotter")
)

func main() {
	flag.Parse()
	lvl, err := logrus.ParseLevel(*logLevel)
	if err != nil {
		log.L.WithError(err).Fatal("failed to prepare logger")
	}
	logrus.SetLevel(lvl)
	logrus.SetFormatter(&logrus.JSONFormatter{
		TimestampFormat: log.RFC3339NanoFixed,
	})

	var (
		ctx    = log.WithLogger(context.Background(), log.L)
		config Config
	)

	// Get configuration from specified file
	if *configPath != "" {
		if _, err := toml.DecodeFile(*configPath, &config); err != nil {
			log.G(ctx).WithError(err).Fatalf("failed to load config file %q", *configPath)
		}
	}

	// Prepare kubeconfig-based keychain if required
	kc := authn.DefaultKeychain
	if config.KubeconfigKeychainConfig.EnableKeychain {
		var opts []keychain.KubeconfigOption
		if kcp := config.KubeconfigKeychainConfig.KubeconfigPath; kcp != "" {
			opts = append(opts, keychain.WithKubeconfigPath(kcp))
		}
		kc = authn.NewMultiKeychain(kc, keychain.NewKubeconfigKeychain(ctx, opts...))
	}

	// Use RegistryHosts based on ResolverConfig and keychain
	hosts := hostsFromConfig(config.ResolverConfig, kc)

	// Configure filesystem and snapshotter
	fs, err := stargzfs.NewFilesystem(filepath.Join(*rootDir, "stargz"),
		config.Config,
		stargzfs.WithGetSources(sources(
			sourceFromCRILabels(hosts),      // provides source info based on CRI labels
			source.FromDefaultLabels(hosts), // provides source info based on default labels
		)),
	)
	if err != nil {
		log.G(ctx).WithError(err).Fatalf("failed to configure filesystem")
	}
	rs, err := snbase.NewSnapshotter(ctx, filepath.Join(*rootDir, "snapshotter"), fs, snbase.AsynchronousRemove)
	if err != nil {
		log.G(ctx).WithError(err).Fatalf("failed to configure snapshotter")
	}
	defer func() {
		log.G(ctx).Debug("Closing the snapshotter")
		rs.Close()
		log.G(ctx).Info("Exiting")
	}()

	// Create a gRPC server
	rpc := grpc.NewServer()

	// Convert the snapshotter to a gRPC service,
	service := snapshotservice.FromSnapshotter(rs)

	// Register the service with the gRPC server
	snapshotsapi.RegisterSnapshotsServer(rpc, service)

	// Prepare the directory for the socket
	if err := os.MkdirAll(filepath.Dir(*address), 0700); err != nil {
		log.G(ctx).WithError(err).Fatalf("failed to create directory %q", filepath.Dir(*address))
	}

	// Try to remove the socket file to avoid EADDRINUSE
	if err := os.RemoveAll(*address); err != nil {
		log.G(ctx).WithError(err).Fatalf("failed to remove %q", *address)
	}

	// Listen and serve
	l, err := net.Listen("unix", *address)
	if err != nil {
		log.G(ctx).WithError(err).Fatalf("error on listen socket %q", *address)
	}
	go func() {
		if err := rpc.Serve(l); err != nil {
			log.G(ctx).WithError(err).Fatalf("error on serving via socket %q", *address)
		}
	}()
	waitForSIGINT()
	log.G(ctx).Info("Got SIGINT")
}

func waitForSIGINT() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c
}

func hostsFromConfig(cfg ResolverConfig, keychain authn.Keychain) docker.RegistryHosts {
	return func(host string) (hosts []docker.RegistryHost, _ error) {
		for _, h := range append(cfg.Host[host].Mirrors, MirrorConfig{
			Host: host,
		}) {
			tr := &http.Client{Transport: http.DefaultTransport.(*http.Transport).Clone()}
			config := docker.RegistryHost{
				Client:       tr,
				Host:         h.Host,
				Scheme:       "https",
				Path:         "/v2",
				Capabilities: docker.HostCapabilityPull,
				Authorizer: docker.NewDockerAuthorizer(
					docker.WithAuthClient(tr),
					docker.WithAuthCreds(func(host string) (string, string, error) {
						if host == "registry-1.docker.io" {
							host = "index.docker.io"
						}
						reg, err := name.NewRegistry(host)
						if err != nil {
							return "", "", err
						}
						authn, err := keychain.Resolve(reg)
						if err != nil {
							return "", "", err
						}
						acfg, err := authn.Authorization()
						if err != nil {
							return "", "", err
						}
						if acfg.IdentityToken != "" {
							return "", acfg.IdentityToken, nil
						}
						return acfg.Username, acfg.Password, nil
					})),
			}
			if localhost, _ := docker.MatchLocalhost(config.Host); localhost || h.Insecure {
				config.Scheme = "http"
			}
			if config.Host == "docker.io" {
				config.Host = "registry-1.docker.io"
			}
			hosts = append(hosts, config)
		}
		return
	}
}

func sources(ps ...source.GetSources) source.GetSources {
	return func(labels map[string]string) (source []source.Source, allErr error) {
		for _, p := range ps {
			src, err := p(labels)
			if err == nil {
				return src, nil
			}
			allErr = multierror.Append(allErr, err)
		}
		return
	}
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"
	fixtures "gopkg.in/src-d/go-git-fixtures.v3"
)

func TestDirectories(t *testing.T) {
	tests := []struct {
		path     string
		expected directory
		error    bool
	}{
		{
			path:     "relative",
			expected: directory{Path: "relative"},
		},
		{
			path:     "longer/relative",
			expected: directory{Path: "longer/relative"},
		},
		{
			path:     "/absolute",
			expected: directory{Path: "/absolute"},
		},
		{
			path:     "/longer/absolute",
			expected: directory{Path: "/longer/absolute"},
		},
		{
			path:     "file://relative",
			expected: directory{Path: "relative"},
		},
		{
			path:     "file://longer/relative",
			expected: directory{Path: "longer/relative"},
		},
		{
			path:     "file:///absolute",
			expected: directory{Path: "/absolute"},
		},
		{
			path:     "file:///longer/absolute",
			expected: directory{Path: "/longer/absolute"},
		},
		{
			path:  "http://relative",
			error: true,
		},
		{
			path: "file:///siva/path?format=siva",
			expected: directory{
				Path:   "/siva/path",
				Format: "siva",
			},
		},
		{
			path: "file:///siva/path?format=git",
			expected: directory{
				Path:   "/siva/path",
				Format: "git",
			},
		},
		{
			path:  "file:///siva/path?format=nope",
			error: true,
		},
		{
			path: "file:///siva/path?bare=true",
			expected: directory{
				Path: "/siva/path",
				Bare: bareOn,
			},
		},
		{
			path: "file:///siva/path?bare=false",
			expected: directory{
				Path: "/siva/path",
				Bare: bareOff,
			},
		},
		{
			path:  "file:///siva/path?bare=nope",
			error: true,
		},
		{
			path: "file:///siva/path?rooted=true",
			expected: directory{
				Path:   "/siva/path",
				Rooted: true,
			},
		},
		{
			path: "file:///siva/path?rooted=false",
			expected: directory{
				Path:   "/siva/path",
				Rooted: false,
			},
		},
		{
			path:  "file:///siva/path?bare=nope",
			error: true,
		},
		{
			path: "file:///siva/path?bucket=42",
			expected: directory{
				Path:   "/siva/path",
				Bucket: 42,
			},
		},
		{
			path:  "file:///siva/path?bucket=false",
			error: true,
		},
		{
			path: "file:///siva/path?format=git&bare=false",
			expected: directory{
				Path:   "/siva/path",
				Format: "git",
				Bare:   bareOff,
			},
		},
		{
			path: "file:///siva/path?format=siva&rooted=false&bucket=42",
			expected: directory{
				Path:   "/siva/path",
				Format: "siva",
				Rooted: false,
				Bucket: 42,
			},
		},
	}

	for _, test := range tests {
		t.Run(test.path, func(t *testing.T) {
			require := require.New(t)
			dir := directory{Path: test.path}

			dir, err := parseDirectory(dir)
			if test.error {
				require.Error(err)
				return
			}

			require.Equal(test.expected, dir)
		})
	}
}

func TestDiscoverBare(t *testing.T) {
	err := fixtures.Init()
	require.NoError(t, err)
	defer func() {
		_ = fixtures.Clean()
	}

	tmpDir, err := ioutil.TempDir("", "gitbase")
	require.NoError(t, err)
	defer os.RemoveAll(tmpDir)

	emptyDir := filepath.Join(tmpDir, "empty")
	err = os.Mkdir(emptyDir, 0777)
	require.NoError(t, err)

	bareDir := filepath.Join(tmpDir, "bare")
	err = os.Mkdir(bareDir, 0777)
	require.NoError(t, err)
	dir := fixtures.ByTag("worktree").One().DotGit().Root()
	err = os.Rename(dir, filepath.Join(bareDir, "repo"))
	require.NoError(t, err)

	nonBareDir := filepath.Join(tmpDir, "non_bare")
	err = os.Mkdir(nonBareDir, 0777)
	require.NoError(t, err)
	dir = fixtures.ByTag("worktree").One().Worktree().Root()
	err = os.Rename(dir, filepath.Join(nonBareDir, "repo"))
	require.NoError(t, err)

	tests := []struct {
		path     string
		bare     bareness
		expected bool
		err      bool
	}{
		{
			path: "/does/not/exist",
			err:  true,
		},
		{
			path:     emptyDir,
			bare:     bareAuto,
			expected: false,
		},
		{
			path:     emptyDir,
			bare:     bareOn,
			expected: true,
		},
		{
			path:     emptyDir,
			bare:     bareOff,
			expected: false,
		},
		{
			path:     bareDir,
			bare:     bareAuto,
			expected: true,
		},
		{
			path:     bareDir,
			bare:     bareOn,
			expected: true,
		},
		{
			path:     bareDir,
			bare:     bareOff,
			expected: false,
		},
		{
			path:     nonBareDir,
			bare:     bareAuto,
			expected: false,
		},
		{
			path:     nonBareDir,
			bare:     bareOn,
			expected: true,
		},
		{
			path:     nonBareDir,
			bare:     bareOff,
			expected: false,
		},
	}

	for _, test := range tests {
		dir := directory{
			Path: test.path,
			Bare: test.bare,
		}

		t.Run(bareTestName(dir, test.err), func(t *testing.T) {
			bare, err := discoverBare(dir)
			if test.err {
				require.Error(t, err)
				return
			}

			require.NoError(t, err)
			require.Equal(t, test.expected, bare)
		})
	}
}

func bareTestName(d directory, err bool) string {
	bare := ""
	switch d.Bare {
	case bareOn:
		bare = "bare"
	case bareOff:
		bare = "non-bare"
	case bareAuto:
		bare = "auto"
	}

	if err {
		bare = "error"
	}

	return fmt.Sprintf("%s_%s", d.Path, bare)
}
<commit_msg>cli: fix server test<commit_after>package command

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	fixtures "github.com/src-d/go-git-fixtures"
	"github.com/stretchr/testify/require"
)

func TestDirectories(t *testing.T) {
	tests := []struct {
		path     string
		expected directory
		error    bool
	}{
		{
			path:     "relative",
			expected: directory{Path: "relative"},
		},
		{
			path:     "longer/relative",
			expected: directory{Path: "longer/relative"},
		},
		{
			path:     "/absolute",
			expected: directory{Path: "/absolute"},
		},
		{
			path:     "/longer/absolute",
			expected: directory{Path: "/longer/absolute"},
		},
		{
			path:     "file://relative",
			expected: directory{Path: "relative"},
		},
		{
			path:     "file://longer/relative",
			expected: directory{Path: "longer/relative"},
		},
		{
			path:     "file:///absolute",
			expected: directory{Path: "/absolute"},
		},
		{
			path:     "file:///longer/absolute",
			expected: directory{Path: "/longer/absolute"},
		},
		{
			path:  "http://relative",
			error: true,
		},
		{
			path: "file:///siva/path?format=siva",
			expected: directory{
				Path:   "/siva/path",
				Format: "siva",
			},
		},
		{
			path: "file:///siva/path?format=git",
			expected: directory{
				Path:   "/siva/path",
				Format: "git",
			},
		},
		{
			path:  "file:///siva/path?format=nope",
			error: true,
		},
		{
			path: "file:///siva/path?bare=true",
			expected: directory{
				Path: "/siva/path",
				Bare: bareOn,
			},
		},
		{
			path: "file:///siva/path?bare=false",
			expected: directory{
				Path: "/siva/path",
				Bare: bareOff,
			},
		},
		{
			path:  "file:///siva/path?bare=nope",
			error: true,
		},
		{
			path: "file:///siva/path?rooted=true",
			expected: directory{
				Path:   "/siva/path",
				Rooted: true,
			},
		},
		{
			path: "file:///siva/path?rooted=false",
			expected: directory{
				Path:   "/siva/path",
				Rooted: false,
			},
		},
		{
			path:  "file:///siva/path?bare=nope",
			error: true,
		},
		{
			path: "file:///siva/path?bucket=42",
			expected: directory{
				Path:   "/siva/path",
				Bucket: 42,
			},
		},
		{
			path:  "file:///siva/path?bucket=false",
			error: true,
		},
		{
			path: "file:///siva/path?format=git&bare=false",
			expected: directory{
				Path:   "/siva/path",
				Format: "git",
				Bare:   bareOff,
			},
		},
		{
			path: "file:///siva/path?format=siva&rooted=false&bucket=42",
			expected: directory{
				Path:   "/siva/path",
				Format: "siva",
				Rooted: false,
				Bucket: 42,
			},
		},
	}

	for _, test := range tests {
		t.Run(test.path, func(t *testing.T) {
			require := require.New(t)
			dir := directory{Path: test.path}

			dir, err := parseDirectory(dir)
			if test.error {
				require.Error(err)
				return
			}

			require.Equal(test.expected, dir)
		})
	}
}

func TestDiscoverBare(t *testing.T) {
	defer func() {
		require.NoError(t, fixtures.Clean())
	}()

	tmpDir, err := ioutil.TempDir("", "gitbase")
	require.NoError(t, err)
	defer os.RemoveAll(tmpDir)

	emptyDir := filepath.Join(tmpDir, "empty")
	err = os.Mkdir(emptyDir, 0777)
	require.NoError(t, err)

	bareDir := filepath.Join(tmpDir, "bare")
	err = os.Mkdir(bareDir, 0777)
	require.NoError(t, err)
	dir := fixtures.ByTag("worktree").One().DotGit().Root()
	err = os.Rename(dir, filepath.Join(bareDir, "repo"))
	require.NoError(t, err)

	nonBareDir := filepath.Join(tmpDir, "non_bare")
	err = os.Mkdir(nonBareDir, 0777)
	require.NoError(t, err)
	dir = fixtures.ByTag("worktree").One().Worktree().Root()
	err = os.Rename(dir, filepath.Join(nonBareDir, "repo"))
	require.NoError(t, err)

	tests := []struct {
		path     string
		bare     bareness
		expected bool
		err      bool
	}{
		{
			path: "/does/not/exist",
			err:  true,
		},
		{
			path:     emptyDir,
			bare:     bareAuto,
			expected: false,
		},
		{
			path:     emptyDir,
			bare:     bareOn,
			expected: true,
		},
		{
			path:     emptyDir,
			bare:     bareOff,
			expected: false,
		},
		{
			path:     bareDir,
			bare:     bareAuto,
			expected: true,
		},
		{
			path:     bareDir,
			bare:     bareOn,
			expected: true,
		},
		{
			path:     bareDir,
			bare:     bareOff,
			expected: false,
		},
		{
			path:     nonBareDir,
			bare:     bareAuto,
			expected: false,
		},
		{
			path:     nonBareDir,
			bare:     bareOn,
			expected: true,
		},
		{
			path:     nonBareDir,
			bare:     bareOff,
			expected: false,
		},
	}

	for _, test := range tests {
		dir := directory{
			Path: test.path,
			Bare: test.bare,
		}

		t.Run(bareTestName(dir, test.err), func(t *testing.T) {
			bare, err := discoverBare(dir)
			if test.err {
				require.Error(t, err)
				return
			}

			require.NoError(t, err)
			require.Equal(t, test.expected, bare)
		})
	}
}

func bareTestName(d directory, err bool) string {
	bare := ""
	switch d.Bare {
	case bareOn:
		bare = "bare"
	case bareOff:
		bare = "non-bare"
	case bareAuto:
		bare = "auto"
	}

	if err {
		bare = "error"
	}

	return fmt.Sprintf("%s_%s", d.Path, bare)
}
<|endoftext|>


<commit_before>package collectors

import (
	"fmt"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"golang.org/x/net/context"
	"google.golang.org/api/monitoring/v3"

	"github.com/frodenas/stackdriver_exporter/utils"
)

type MonitoringCollector struct {
	projectID                       string
	metricsPrefix                   string
	metricsInterval                 time.Duration
	monitoringService               *monitoring.Service
	apiCallsTotalMetric             prometheus.Counter
	scrapesTotalMetric              prometheus.Counter
	scrapeErrorsTotalMetric         prometheus.Counter
	lastScrapeErrorMetric           prometheus.Gauge
	lastScrapeTimestampMetric       prometheus.Gauge
	lastScrapeDurationSecondsMetric prometheus.Gauge
}

func NewMonitoringCollector(projectID string, metricsPrefix string, metricsInterval time.Duration, monitoringService *monitoring.Service) (*MonitoringCollector, error) {
	apiCallsTotalMetric := prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "api_calls_total",
			Help:        "Total number of Google Stackdriver Monitoring API calls made.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	scrapesTotalMetric := prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "scrapes_total",
			Help:        "Total number of Google Stackdriver Monitoring metrics scrapes.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	scrapeErrorsTotalMetric := prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "scrape_errors_total",
			Help:        "Total number of Google Stackdriver Monitoring metrics scrape errors.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	lastScrapeErrorMetric := prometheus.NewGauge(
		prometheus.GaugeOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "last_scrape_error",
			Help:        "Whether the last metrics scrape from Google Stackdriver Monitoring resulted in an error (1 for error, 0 for success).",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	lastScrapeTimestampMetric := prometheus.NewGauge(
		prometheus.GaugeOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "last_scrape_timestamp",
			Help:        "Number of seconds since 1970 since last metrics scrape from Google Stackdriver Monitoring.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	lastScrapeDurationSecondsMetric := prometheus.NewGauge(
		prometheus.GaugeOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "last_scrape_duration_seconds",
			Help:        "Duration of the last metrics scrape from Google Stackdriver Monitoring.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	monitoringCollector := &MonitoringCollector{
		projectID:                       projectID,
		metricsPrefix:                   metricsPrefix,
		metricsInterval:                 metricsInterval,
		monitoringService:               monitoringService,
		apiCallsTotalMetric:             apiCallsTotalMetric,
		scrapesTotalMetric:              scrapesTotalMetric,
		scrapeErrorsTotalMetric:         scrapeErrorsTotalMetric,
		lastScrapeErrorMetric:           lastScrapeErrorMetric,
		lastScrapeTimestampMetric:       lastScrapeTimestampMetric,
		lastScrapeDurationSecondsMetric: lastScrapeDurationSecondsMetric,
	}

	return monitoringCollector, nil
}

func (c *MonitoringCollector) Describe(ch chan<- *prometheus.Desc) {
	c.apiCallsTotalMetric.Describe(ch)
	c.scrapesTotalMetric.Describe(ch)
	c.scrapeErrorsTotalMetric.Describe(ch)
	c.lastScrapeErrorMetric.Describe(ch)
	c.lastScrapeTimestampMetric.Describe(ch)
	c.lastScrapeDurationSecondsMetric.Describe(ch)
}

func (c *MonitoringCollector) Collect(ch chan<- prometheus.Metric) {
	var begun = time.Now()

	errorMetric := float64(0)
	if err := c.reportMonitoringMetrics(ch); err != nil {
		errorMetric = float64(1)
		c.scrapeErrorsTotalMetric.Inc()
		log.Errorf("Error while getting Google Stackdriver Monitoring metrics: %s", err)
	}
	c.scrapeErrorsTotalMetric.Collect(ch)

	c.apiCallsTotalMetric.Collect(ch)

	c.scrapesTotalMetric.Inc()
	c.scrapesTotalMetric.Collect(ch)

	c.lastScrapeErrorMetric.Set(errorMetric)
	c.lastScrapeErrorMetric.Collect(ch)

	c.lastScrapeTimestampMetric.Set(float64(time.Now().Unix()))
	c.lastScrapeTimestampMetric.Collect(ch)

	c.lastScrapeDurationSecondsMetric.Set(time.Since(begun).Seconds())
	c.lastScrapeDurationSecondsMetric.Collect(ch)
}

func (c *MonitoringCollector) reportMonitoringMetrics(ch chan<- prometheus.Metric) error {
	metricDescriptorsFunction := func(page *monitoring.ListMetricDescriptorsResponse) error {
		var wg = &sync.WaitGroup{}

		c.apiCallsTotalMetric.Inc()

		doneChannel := make(chan bool, 1)
		errChannel := make(chan error, 1)

		startTime := time.Now().UTC().Add(c.metricsInterval * -1)
		endTime := time.Now().UTC()

		for _, metricDescriptor := range page.MetricDescriptors {
			wg.Add(1)
			go func(metricDescriptor *monitoring.MetricDescriptor, ch chan<- prometheus.Metric) {
				defer wg.Done()
				log.Debugf("Retrieving Google Stackdriver Monitoring metrics for descriptor `%s`...", metricDescriptor.Type)
				timeSeriesListCall := c.monitoringService.Projects.TimeSeries.List(utils.ProjectResource(c.projectID)).
					Filter(fmt.Sprintf("metric.type=\"%s\"", metricDescriptor.Type)).
					IntervalStartTime(startTime.Format(time.RFC3339Nano)).
					IntervalEndTime(endTime.Format(time.RFC3339Nano))

				for {
					c.apiCallsTotalMetric.Inc()
					page, err := timeSeriesListCall.Do()
					if err != nil {
						errChannel <- err
					}
					if err := c.reportTimeSeriesMetrics(page, metricDescriptor, ch); err != nil {
						errChannel <- err
					}
					if page.NextPageToken == "" {
						break
					}
					timeSeriesListCall.PageToken(page.NextPageToken)
				}
			}(metricDescriptor, ch)
		}

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

		select {
		case <-doneChannel:
		case err := <-errChannel:
			return err
		}

		return nil
	}

	log.Debugf("Listing Google Stackdriver Monitoring metric descriptors starting with `%s`...", c.metricsPrefix)
	ctx := context.Background()
	if err := c.monitoringService.Projects.MetricDescriptors.List(utils.ProjectResource(c.projectID)).
		Filter(fmt.Sprintf("metric.type = starts_with(\"%s\")", c.metricsPrefix)).
		Pages(ctx, metricDescriptorsFunction); err != nil {
		return err
	}

	return nil
}

func (c *MonitoringCollector) reportTimeSeriesMetrics(page *monitoring.ListTimeSeriesResponse, metricDescriptor *monitoring.MetricDescriptor, ch chan<- prometheus.Metric) error {
	var metricValue float64
	var metricValueType prometheus.ValueType

	for _, timeSeries := range page.TimeSeries {
		if len(timeSeries.Points) < 1 {
			continue
		}

		switch timeSeries.MetricKind {
		case "GAUGE":
			metricValueType = prometheus.GaugeValue
		case "DELTA":
			metricValueType = prometheus.CounterValue
		case "CUMULATIVE":
			metricValueType = prometheus.CounterValue
		default:
			continue
		}

		switch timeSeries.ValueType {
		case "BOOL":
			metricValue = 0
			if *timeSeries.Points[0].Value.BoolValue {
				metricValue = 1
			}
		case "INT64":
			metricValue = float64(*timeSeries.Points[0].Value.Int64Value)
		case "DOUBLE":
			metricValue = *timeSeries.Points[0].Value.DoubleValue
		default:
			log.Debugf("Discarding `%s` metric: %+v", timeSeries.ValueType, timeSeries)
			continue
		}

		labelKeys := []string{"unit", "resource_type"}
		labelValues := []string{metricDescriptor.Unit, timeSeries.Resource.Type}
		for key, value := range timeSeries.Metric.Labels {
			labelKeys = append(labelKeys, key)
			labelValues = append(labelValues, value)
		}
		for key, value := range timeSeries.Resource.Labels {
			labelKeys = append(labelKeys, key)
			labelValues = append(labelValues, value)
		}

		ch <- prometheus.MustNewConstMetric(
			prometheus.NewDesc(
				prometheus.BuildFQName("stackdriver", "monitoring", utils.NormalizeMetricName(timeSeries.Metric.Type)),
				metricDescriptor.Description,
				labelKeys,
				prometheus.Labels{},
			),
			metricValueType,
			metricValue,
			labelValues...,
		)
	}

	return nil
}
<commit_msg>Collect only the newest timeseries point<commit_after>package collectors

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"golang.org/x/net/context"
	"google.golang.org/api/monitoring/v3"

	"github.com/frodenas/stackdriver_exporter/utils"
)

type MonitoringCollector struct {
	projectID                       string
	metricsPrefix                   string
	metricsInterval                 time.Duration
	monitoringService               *monitoring.Service
	apiCallsTotalMetric             prometheus.Counter
	scrapesTotalMetric              prometheus.Counter
	scrapeErrorsTotalMetric         prometheus.Counter
	lastScrapeErrorMetric           prometheus.Gauge
	lastScrapeTimestampMetric       prometheus.Gauge
	lastScrapeDurationSecondsMetric prometheus.Gauge
}

func NewMonitoringCollector(projectID string, metricsPrefix string, metricsInterval time.Duration, monitoringService *monitoring.Service) (*MonitoringCollector, error) {
	apiCallsTotalMetric := prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "api_calls_total",
			Help:        "Total number of Google Stackdriver Monitoring API calls made.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	scrapesTotalMetric := prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "scrapes_total",
			Help:        "Total number of Google Stackdriver Monitoring metrics scrapes.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	scrapeErrorsTotalMetric := prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "scrape_errors_total",
			Help:        "Total number of Google Stackdriver Monitoring metrics scrape errors.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	lastScrapeErrorMetric := prometheus.NewGauge(
		prometheus.GaugeOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "last_scrape_error",
			Help:        "Whether the last metrics scrape from Google Stackdriver Monitoring resulted in an error (1 for error, 0 for success).",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	lastScrapeTimestampMetric := prometheus.NewGauge(
		prometheus.GaugeOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "last_scrape_timestamp",
			Help:        "Number of seconds since 1970 since last metrics scrape from Google Stackdriver Monitoring.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	lastScrapeDurationSecondsMetric := prometheus.NewGauge(
		prometheus.GaugeOpts{
			Namespace:   "stackdriver",
			Subsystem:   "monitoring",
			Name:        "last_scrape_duration_seconds",
			Help:        "Duration of the last metrics scrape from Google Stackdriver Monitoring.",
			ConstLabels: prometheus.Labels{"project_id": projectID},
		},
	)

	monitoringCollector := &MonitoringCollector{
		projectID:                       projectID,
		metricsPrefix:                   metricsPrefix,
		metricsInterval:                 metricsInterval,
		monitoringService:               monitoringService,
		apiCallsTotalMetric:             apiCallsTotalMetric,
		scrapesTotalMetric:              scrapesTotalMetric,
		scrapeErrorsTotalMetric:         scrapeErrorsTotalMetric,
		lastScrapeErrorMetric:           lastScrapeErrorMetric,
		lastScrapeTimestampMetric:       lastScrapeTimestampMetric,
		lastScrapeDurationSecondsMetric: lastScrapeDurationSecondsMetric,
	}

	return monitoringCollector, nil
}

func (c *MonitoringCollector) Describe(ch chan<- *prometheus.Desc) {
	c.apiCallsTotalMetric.Describe(ch)
	c.scrapesTotalMetric.Describe(ch)
	c.scrapeErrorsTotalMetric.Describe(ch)
	c.lastScrapeErrorMetric.Describe(ch)
	c.lastScrapeTimestampMetric.Describe(ch)
	c.lastScrapeDurationSecondsMetric.Describe(ch)
}

func (c *MonitoringCollector) Collect(ch chan<- prometheus.Metric) {
	var begun = time.Now()

	errorMetric := float64(0)
	if err := c.reportMonitoringMetrics(ch); err != nil {
		errorMetric = float64(1)
		c.scrapeErrorsTotalMetric.Inc()
		log.Errorf("Error while getting Google Stackdriver Monitoring metrics: %s", err)
	}
	c.scrapeErrorsTotalMetric.Collect(ch)

	c.apiCallsTotalMetric.Collect(ch)

	c.scrapesTotalMetric.Inc()
	c.scrapesTotalMetric.Collect(ch)

	c.lastScrapeErrorMetric.Set(errorMetric)
	c.lastScrapeErrorMetric.Collect(ch)

	c.lastScrapeTimestampMetric.Set(float64(time.Now().Unix()))
	c.lastScrapeTimestampMetric.Collect(ch)

	c.lastScrapeDurationSecondsMetric.Set(time.Since(begun).Seconds())
	c.lastScrapeDurationSecondsMetric.Collect(ch)
}

func (c *MonitoringCollector) reportMonitoringMetrics(ch chan<- prometheus.Metric) error {
	metricDescriptorsFunction := func(page *monitoring.ListMetricDescriptorsResponse) error {
		var wg = &sync.WaitGroup{}

		c.apiCallsTotalMetric.Inc()

		doneChannel := make(chan bool, 1)
		errChannel := make(chan error, 1)

		startTime := time.Now().UTC().Add(c.metricsInterval * -1)
		endTime := time.Now().UTC()

		for _, metricDescriptor := range page.MetricDescriptors {
			wg.Add(1)
			go func(metricDescriptor *monitoring.MetricDescriptor, ch chan<- prometheus.Metric) {
				defer wg.Done()
				log.Debugf("Retrieving Google Stackdriver Monitoring metrics for descriptor `%s`...", metricDescriptor.Type)
				timeSeriesListCall := c.monitoringService.Projects.TimeSeries.List(utils.ProjectResource(c.projectID)).
					Filter(fmt.Sprintf("metric.type=\"%s\"", metricDescriptor.Type)).
					IntervalStartTime(startTime.Format(time.RFC3339Nano)).
					IntervalEndTime(endTime.Format(time.RFC3339Nano))

				for {
					c.apiCallsTotalMetric.Inc()
					page, err := timeSeriesListCall.Do()
					if err != nil {
						errChannel <- err
					}
					if err := c.reportTimeSeriesMetrics(page, metricDescriptor, ch); err != nil {
						errChannel <- err
					}
					if page.NextPageToken == "" {
						break
					}
					timeSeriesListCall.PageToken(page.NextPageToken)
				}
			}(metricDescriptor, ch)
		}

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

		select {
		case <-doneChannel:
		case err := <-errChannel:
			return err
		}

		return nil
	}

	log.Debugf("Listing Google Stackdriver Monitoring metric descriptors starting with `%s`...", c.metricsPrefix)
	ctx := context.Background()
	if err := c.monitoringService.Projects.MetricDescriptors.List(utils.ProjectResource(c.projectID)).
		Filter(fmt.Sprintf("metric.type = starts_with(\"%s\")", c.metricsPrefix)).
		Pages(ctx, metricDescriptorsFunction); err != nil {
		return err
	}

	return nil
}

func (c *MonitoringCollector) reportTimeSeriesMetrics(page *monitoring.ListTimeSeriesResponse, metricDescriptor *monitoring.MetricDescriptor, ch chan<- prometheus.Metric) error {
	var metricValue float64
	var metricValueType prometheus.ValueType
	var newestTSPoint *monitoring.Point

	for _, timeSeries := range page.TimeSeries {
		newestEndTime := time.Unix(0, 0)
		for _, point := range timeSeries.Points {
			endTime, err := time.Parse(time.RFC3339Nano, point.Interval.EndTime)
			if err != nil {
				return errors.New(fmt.Sprintf("Error parsing TimeSeries Point interval end time `%s`: %s", point.Interval.EndTime, err))
			}
			if endTime.After(newestEndTime) {
				newestEndTime = endTime
				newestTSPoint = point
			}
		}

		switch timeSeries.MetricKind {
		case "GAUGE":
			metricValueType = prometheus.GaugeValue
		case "DELTA":
			metricValueType = prometheus.CounterValue
		case "CUMULATIVE":
			metricValueType = prometheus.CounterValue
		default:
			continue
		}

		switch timeSeries.ValueType {
		case "BOOL":
			metricValue = 0
			if *newestTSPoint.Value.BoolValue {
				metricValue = 1
			}
		case "INT64":
			metricValue = float64(*newestTSPoint.Value.Int64Value)
		case "DOUBLE":
			metricValue = *newestTSPoint.Value.DoubleValue
		default:
			log.Debugf("Discarding `%s` metric: %+v", timeSeries.ValueType, timeSeries)
			continue
		}

		labelKeys := []string{"unit", "resource_type"}
		labelValues := []string{metricDescriptor.Unit, timeSeries.Resource.Type}
		for key, value := range timeSeries.Metric.Labels {
			labelKeys = append(labelKeys, key)
			labelValues = append(labelValues, value)
		}
		for key, value := range timeSeries.Resource.Labels {
			labelKeys = append(labelKeys, key)
			labelValues = append(labelValues, value)
		}

		ch <- prometheus.MustNewConstMetric(
			prometheus.NewDesc(
				prometheus.BuildFQName("stackdriver", "monitoring", utils.NormalizeMetricName(timeSeries.Metric.Type)),
				metricDescriptor.Description,
				labelKeys,
				prometheus.Labels{},
			),
			metricValueType,
			metricValue,
			labelValues...,
		)
	}

	return nil
}
<|endoftext|>


<commit_before>package std

import (
	"math"

	"github.com/DeedleFake/wdte"
)

func save(f wdte.Func, saved ...wdte.Func) wdte.Func {
	return wdte.GoFunc(func(frame wdte.Frame, args ...wdte.Func) wdte.Func {
		return f.Call(frame, append(saved, args...)...)
	})
}

// Add returns the sum of its arguments. If called with only 1
// argument, it returns a function which adds arguments given to that
// one argument.
func Add(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Add)

	case 1:
		return save(wdte.GoFunc(Add), args[0])
	}

	frame = frame.WithID("+")

	var sum wdte.Number
	for _, arg := range args {
		arg = arg.Call(frame)
		if _, ok := arg.(error); ok {
			return arg
		}
		sum += arg.(wdte.Number)
	}
	return sum
}

// Sub returns args[0] - args[1]. If called with only 1 argument, it
// returns a function which returns that argument minus the argument
// given.
func Sub(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Sub)

	case 1:
		return save(wdte.GoFunc(Sub), args[0])
	}

	frame = frame.WithID("-")

	a1 := args[0].Call(frame)
	if _, ok := a1.(error); ok {
		return a1
	}

	a2 := args[1].Call(frame)
	if _, ok := a2.(error); ok {
		return a2
	}

	return a1.(wdte.Number) - a2.(wdte.Number)
}

// Mult returns the product of its arguments. If called with only 1
// argument, it returns a function that multiplies that argument by
// its own arguments.
func Mult(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Mult)

	case 1:
		return save(wdte.GoFunc(Mult), args[0])
	}

	frame = frame.WithID("*")

	p := wdte.Number(1)
	for _, arg := range args {
		arg = arg.Call(frame)
		if _, ok := arg.(error); ok {
			return arg
		}
		p *= arg.(wdte.Number)
	}
	return p
}

// Div returns args[0] / args[1]. If called with only 1 argument, it
// returns a function which divides its own argument by the original
// argument.
func Div(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Div)

	case 1:
		return save(wdte.GoFunc(Div), args[0])
	}

	frame = frame.WithID("/")

	a1 := args[0].Call(frame)
	if _, ok := a1.(error); ok {
		return a1
	}

	a2 := args[1].Call(frame)
	if _, ok := a2.(error); ok {
		return a2
	}

	return a1.(wdte.Number) / a2.(wdte.Number)
}

// Mod returns args[0] % args[1]. If called with only 1 argument, it
// returns a function which divides its own argument by the original
// argument.
func Mod(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Mod)

	case 1:
		return save(wdte.GoFunc(Mod), args[0])
	}

	frame = frame.WithID("%")

	a1 := args[0].Call(frame)
	if _, ok := a1.(error); ok {
		return a1
	}

	a2 := args[1].Call(frame)
	if _, ok := a2.(error); ok {
		return a2
	}

	return wdte.Number(math.Mod(
		float64(a1.(wdte.Number)),
		float64(a2.(wdte.Number)),
	))
}

func Equals(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Equals)

	case 1:
		return save(wdte.GoFunc(Equals), args[0])
	}

	a1 := args[0].Call(frame)
	if _, ok := a1.(error); ok {
		return a1
	}

	a2 := args[1].Call(frame)
	if _, ok := a2.(error); ok {
		return a2
	}

	if cmp, ok := a1.(wdte.Comparer); ok {
		c, _ := cmp.Compare(a2)
		return wdte.Bool(c == 0)
	}

	if cmp, ok := a2.(wdte.Comparer); ok {
		c, _ := cmp.Compare(a1)
		return wdte.Bool(c == 0)
	}

	return wdte.Bool(a1 == args[1].Call(frame))
}

// Insert adds the functions in this package to m. It maps
// mathematical functions to the corresponding mathematical symbols.
// For example, Add() becomes `+`, Sub() becomes `-`, and so on.
// Comparisons get mapped to the cooresponding comparison symbols from
// C-style languages. For example, Equals() becomes `==`.
func Insert(m *wdte.Module) {
	m.Funcs["+"] = wdte.GoFunc(Add)
	m.Funcs["-"] = wdte.GoFunc(Sub)
	m.Funcs["*"] = wdte.GoFunc(Mult)
	m.Funcs["/"] = wdte.GoFunc(Div)
	m.Funcs["%"] = wdte.GoFunc(Mod)

	m.Funcs["=="] = wdte.GoFunc(Equals)
}
<commit_msg>std: Add Godoc comment to Equals(). Fix minor problem in the function, too.<commit_after>package std

import (
	"math"

	"github.com/DeedleFake/wdte"
)

func save(f wdte.Func, saved ...wdte.Func) wdte.Func {
	return wdte.GoFunc(func(frame wdte.Frame, args ...wdte.Func) wdte.Func {
		return f.Call(frame, append(saved, args...)...)
	})
}

// Add returns the sum of its arguments. If called with only 1
// argument, it returns a function which adds arguments given to that
// one argument.
func Add(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Add)

	case 1:
		return save(wdte.GoFunc(Add), args[0])
	}

	frame = frame.WithID("+")

	var sum wdte.Number
	for _, arg := range args {
		arg = arg.Call(frame)
		if _, ok := arg.(error); ok {
			return arg
		}
		sum += arg.(wdte.Number)
	}
	return sum
}

// Sub returns args[0] - args[1]. If called with only 1 argument, it
// returns a function which returns that argument minus the argument
// given.
func Sub(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Sub)

	case 1:
		return save(wdte.GoFunc(Sub), args[0])
	}

	frame = frame.WithID("-")

	a1 := args[0].Call(frame)
	if _, ok := a1.(error); ok {
		return a1
	}

	a2 := args[1].Call(frame)
	if _, ok := a2.(error); ok {
		return a2
	}

	return a1.(wdte.Number) - a2.(wdte.Number)
}

// Mult returns the product of its arguments. If called with only 1
// argument, it returns a function that multiplies that argument by
// its own arguments.
func Mult(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Mult)

	case 1:
		return save(wdte.GoFunc(Mult), args[0])
	}

	frame = frame.WithID("*")

	p := wdte.Number(1)
	for _, arg := range args {
		arg = arg.Call(frame)
		if _, ok := arg.(error); ok {
			return arg
		}
		p *= arg.(wdte.Number)
	}
	return p
}

// Div returns args[0] / args[1]. If called with only 1 argument, it
// returns a function which divides its own argument by the original
// argument.
func Div(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Div)

	case 1:
		return save(wdte.GoFunc(Div), args[0])
	}

	frame = frame.WithID("/")

	a1 := args[0].Call(frame)
	if _, ok := a1.(error); ok {
		return a1
	}

	a2 := args[1].Call(frame)
	if _, ok := a2.(error); ok {
		return a2
	}

	return a1.(wdte.Number) / a2.(wdte.Number)
}

// Mod returns args[0] % args[1]. If called with only 1 argument, it
// returns a function which divides its own argument by the original
// argument.
func Mod(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Mod)

	case 1:
		return save(wdte.GoFunc(Mod), args[0])
	}

	frame = frame.WithID("%")

	a1 := args[0].Call(frame)
	if _, ok := a1.(error); ok {
		return a1
	}

	a2 := args[1].Call(frame)
	if _, ok := a2.(error); ok {
		return a2
	}

	return wdte.Number(math.Mod(
		float64(a1.(wdte.Number)),
		float64(a2.(wdte.Number)),
	))
}

// Equals checks if two values are equal or not. If called with only
// one argument, it returns a function which checks that argument for
// equality with other values.
//
// If the first argument given implements wdte.Comparer, it is used
// for the comparison. If not, and the second does, then that is used.
// If neither does, a simple, direct Go equality check is used.
func Equals(frame wdte.Frame, args ...wdte.Func) wdte.Func {
	switch len(args) {
	case 0:
		return wdte.GoFunc(Equals)

	case 1:
		return save(wdte.GoFunc(Equals), args[0])
	}

	a1 := args[0].Call(frame)
	if _, ok := a1.(error); ok {
		return a1
	}

	a2 := args[1].Call(frame)
	if _, ok := a2.(error); ok {
		return a2
	}

	if cmp, ok := a1.(wdte.Comparer); ok {
		c, _ := cmp.Compare(a2)
		return wdte.Bool(c == 0)
	}

	if cmp, ok := a2.(wdte.Comparer); ok {
		c, _ := cmp.Compare(a1)
		return wdte.Bool(c == 0)
	}

	return wdte.Bool(a1 == a2)
}

// Insert adds the functions in this package to m. It maps
// mathematical functions to the corresponding mathematical symbols.
// For example, Add() becomes `+`, Sub() becomes `-`, and so on.
// Comparisons get mapped to the cooresponding comparison symbols from
// C-style languages. For example, Equals() becomes `==`.
func Insert(m *wdte.Module) {
	m.Funcs["+"] = wdte.GoFunc(Add)
	m.Funcs["-"] = wdte.GoFunc(Sub)
	m.Funcs["*"] = wdte.GoFunc(Mult)
	m.Funcs["/"] = wdte.GoFunc(Div)
	m.Funcs["%"] = wdte.GoFunc(Mod)

	m.Funcs["=="] = wdte.GoFunc(Equals)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"image"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
)

const (
	awsKeyEnvVar    = "AWS_ACCESS_KEY_ID"
	awsSecretEnvVar = "AWS_SECRET_ACCESS_KEY"
	s3BucketEnvVar  = "PIXLSERV_S3_BUCKET"
)

var (
	storageImpl storage
)

type storage interface {
	init() error

	loadImage(imagePath string) (image.Image, string, error)

	saveImage(img image.Image, format string, imagePath string) (int, error)

	deleteImage(imagePath string) error

	imageExists(imagePath string) bool
}

func storageInit() error {
	if os.Getenv(awsKeyEnvVar) != "" && os.Getenv(awsSecretEnvVar) != "" && os.Getenv(s3BucketEnvVar) != "" {
		storageImpl = new(s3Storage)
		log.Println("Using S3 storage")
	} else {
		storageImpl = new(localStorage)
		log.Println("Using local storage")
	}

	return storageImpl.init()
}

func storageCleanUp() {
}

func loadImage(imagePath string) (image.Image, string, error) {
	return storageImpl.loadImage(imagePath)
}

func saveImage(img image.Image, format string, imagePath string) (int, error) {
	return storageImpl.saveImage(img, format, imagePath)
}

func deleteImage(imagePath string) error {
	return storageImpl.deleteImage(imagePath)
}

func imageExists(imagePath string) bool {
	return storageImpl.imageExists(imagePath)
}

// localStorage is a storage implementation using local disk
type localStorage struct {
	path string
}

func (s *localStorage) init() error {
	path := os.Getenv(Config.localPath)
	if path == "" {
		path = defaultLocalPath
	}
	s.path = path
	return nil
}

func (s *localStorage) loadImage(imagePath string) (image.Image, string, error) {
	reader, err := os.Open(s.path + "/" + imagePath)
	defer reader.Close()

	if err != nil {
		return nil, "", fmt.Errorf("image not found: %q", imagePath)
	}
	img, format, err := image.Decode(reader)
	if err != nil {
		return nil, "", fmt.Errorf("cannot decode image: %q", imagePath)
	}
	return img, format, nil
}

func (s *localStorage) saveImage(img image.Image, format string, imagePath string) (int, error) {
	// Open file for writing, overwrite if it already exists
	fullPath := s.path + "/" + imagePath
	writer, err := os.Create(fullPath)
	defer writer.Close()

	if err != nil {
		return 0, err
	}

	err = writeImage(img, format, writer)
	if err != nil {
		return 0, err
	}

	f, err := os.Open(fullPath)
	if err != nil {
		return 0, nil
	}

	stat, err := f.Stat()
	if err != nil {
		return 0, nil
	}
	size := stat.Size()

	return int(size), err
}

func (s *localStorage) deleteImage(imagePath string) error {
	return os.Remove(s.path + "/" + imagePath)
}

func (s *localStorage) imageExists(imagePath string) bool {
	if _, err := os.Stat(s.path + "/" + imagePath); os.IsNotExist(err) {
		return false
	}
	return true
}

// s3Storage is a storage implementation using Amazon S3
type s3Storage struct {
	bucket *s3.Bucket
}

func (s *s3Storage) init() error {
	auth, err := aws.EnvAuth()
	if err != nil {
		return err
	}

	bucketName := os.Getenv(s3BucketEnvVar)
	if bucketName == "" {
		return fmt.Errorf("%s not set", s3BucketEnvVar)
	}

	conn := s3.New(auth, aws.EUWest)
	s.bucket = conn.Bucket(bucketName)

	return nil
}

func (s *s3Storage) loadImage(imagePath string) (image.Image, string, error) {
	data, err := s.bucket.Get(imagePath)
	if err != nil {
		return nil, "", err
	}

	format := strings.TrimLeft(filepath.Ext(imagePath), ".")
	image, err := readImage(data, format)
	if err != nil {
		return nil, "", err
	}

	return image, format, nil
}

func (s *s3Storage) saveImage(img image.Image, format string, imagePath string) (int, error) {
	var buffer bytes.Buffer
	err := writeImage(img, format, &buffer)
	if err != nil {
		return 0, err
	}

	size := buffer.Len()
	return size, s.bucket.Put(imagePath, buffer.Bytes(), "image/"+format, s3.PublicRead)
}

func (s *s3Storage) deleteImage(imagePath string) error {
	return s.bucket.Del(imagePath)
}

func (s *s3Storage) imageExists(imagePath string) bool {
	resp, err := s.bucket.List(imagePath, "/", "", 10)
	if err != nil {
		log.Printf("Error while listing S3 bucket: %s\n", err.Error())
		return false
	}
	if resp == nil {
		log.Println("Error while listing S3 bucket: empty response")
	}

	for _, element := range resp.Contents {
		if element.Key == imagePath {
			return true
		}
	}

	return false
}
<commit_msg>Fixed local storage path initialisation<commit_after>package main

import (
	"bytes"
	"fmt"
	"image"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
)

const (
	awsKeyEnvVar    = "AWS_ACCESS_KEY_ID"
	awsSecretEnvVar = "AWS_SECRET_ACCESS_KEY"
	s3BucketEnvVar  = "PIXLSERV_S3_BUCKET"
)

var (
	storageImpl storage
)

type storage interface {
	init() error

	loadImage(imagePath string) (image.Image, string, error)

	saveImage(img image.Image, format string, imagePath string) (int, error)

	deleteImage(imagePath string) error

	imageExists(imagePath string) bool
}

func storageInit() error {
	if os.Getenv(awsKeyEnvVar) != "" && os.Getenv(awsSecretEnvVar) != "" && os.Getenv(s3BucketEnvVar) != "" {
		storageImpl = new(s3Storage)
		log.Println("Using S3 storage")
	} else {
		storageImpl = new(localStorage)
		log.Println("Using local storage")
	}

	return storageImpl.init()
}

func storageCleanUp() {
}

func loadImage(imagePath string) (image.Image, string, error) {
	return storageImpl.loadImage(imagePath)
}

func saveImage(img image.Image, format string, imagePath string) (int, error) {
	return storageImpl.saveImage(img, format, imagePath)
}

func deleteImage(imagePath string) error {
	return storageImpl.deleteImage(imagePath)
}

func imageExists(imagePath string) bool {
	return storageImpl.imageExists(imagePath)
}

// localStorage is a storage implementation using local disk
type localStorage struct {
	path string
}

func (s *localStorage) init() error {
	path := Config.localPath
	s.path = path
	return nil
}

func (s *localStorage) loadImage(imagePath string) (image.Image, string, error) {
	reader, err := os.Open(s.path + "/" + imagePath)
	defer reader.Close()

	if err != nil {
		return nil, "", fmt.Errorf("image not found: %q", imagePath)
	}
	img, format, err := image.Decode(reader)
	if err != nil {
		return nil, "", fmt.Errorf("cannot decode image: %q", imagePath)
	}
	return img, format, nil
}

func (s *localStorage) saveImage(img image.Image, format string, imagePath string) (int, error) {
	// Open file for writing, overwrite if it already exists
	fullPath := s.path + "/" + imagePath
	writer, err := os.Create(fullPath)
	defer writer.Close()

	if err != nil {
		return 0, err
	}

	err = writeImage(img, format, writer)
	if err != nil {
		return 0, err
	}

	f, err := os.Open(fullPath)
	if err != nil {
		return 0, nil
	}

	stat, err := f.Stat()
	if err != nil {
		return 0, nil
	}
	size := stat.Size()

	return int(size), err
}

func (s *localStorage) deleteImage(imagePath string) error {
	return os.Remove(s.path + "/" + imagePath)
}

func (s *localStorage) imageExists(imagePath string) bool {
	if _, err := os.Stat(s.path + "/" + imagePath); os.IsNotExist(err) {
		return false
	}
	return true
}

// s3Storage is a storage implementation using Amazon S3
type s3Storage struct {
	bucket *s3.Bucket
}

func (s *s3Storage) init() error {
	auth, err := aws.EnvAuth()
	if err != nil {
		return err
	}

	bucketName := os.Getenv(s3BucketEnvVar)
	if bucketName == "" {
		return fmt.Errorf("%s not set", s3BucketEnvVar)
	}

	conn := s3.New(auth, aws.EUWest)
	s.bucket = conn.Bucket(bucketName)

	return nil
}

func (s *s3Storage) loadImage(imagePath string) (image.Image, string, error) {
	data, err := s.bucket.Get(imagePath)
	if err != nil {
		return nil, "", err
	}

	format := strings.TrimLeft(filepath.Ext(imagePath), ".")
	image, err := readImage(data, format)
	if err != nil {
		return nil, "", err
	}

	return image, format, nil
}

func (s *s3Storage) saveImage(img image.Image, format string, imagePath string) (int, error) {
	var buffer bytes.Buffer
	err := writeImage(img, format, &buffer)
	if err != nil {
		return 0, err
	}

	size := buffer.Len()
	return size, s.bucket.Put(imagePath, buffer.Bytes(), "image/"+format, s3.PublicRead)
}

func (s *s3Storage) deleteImage(imagePath string) error {
	return s.bucket.Del(imagePath)
}

func (s *s3Storage) imageExists(imagePath string) bool {
	resp, err := s.bucket.List(imagePath, "/", "", 10)
	if err != nil {
		log.Printf("Error while listing S3 bucket: %s\n", err.Error())
		return false
	}
	if resp == nil {
		log.Println("Error while listing S3 bucket: empty response")
	}

	for _, element := range resp.Contents {
		if element.Key == imagePath {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/keybase/kbfs/libkbfs"
	"golang.org/x/net/context"
)

func mdResetOne(
	ctx context.Context, config libkbfs.Config, tlfPath string,
	checkValid, dryRun bool) error {
	handle, err := parseTLFPath(ctx, config.KBPKI(), tlfPath)
	if err != nil {
		return err
	}

	// This function is loosely adapted from
	// folderBranchOps.initMDLocked.

	username, uid, err := config.KBPKI().GetCurrentUserInfo(ctx)
	if err != nil {
		return err
	}

	// Make sure we're a writer before doing anything else.
	if !handle.IsWriter(uid) {
		return libkbfs.NewWriteAccessError(
			handle, username, handle.GetCanonicalPath())
	}

	fmt.Printf("Looking for unmerged branch...\n")

	_, unmergedIRMD, err := config.MDOps().GetForHandle(
		ctx, handle, libkbfs.Unmerged)
	if err != nil {
		return err
	}
	if unmergedIRMD != (libkbfs.ImmutableRootMetadata{}) {
		return fmt.Errorf(
			"%s has unmerged data; try unstaging it first",
			tlfPath)
	}

	fmt.Printf("Getting latest metadata...\n")

	_, irmd, err := config.MDOps().GetForHandle(
		ctx, handle, libkbfs.Merged)
	if err != nil {
		return err
	}

	if irmd == (libkbfs.ImmutableRootMetadata{}) {
		fmt.Printf("No TLF found for %q\n", tlfPath)
		return nil
	}

	if checkValid {
		rootPtr := irmd.Data().Dir.BlockPointer
		if rootPtr.Ref().IsValid() {
			var dirBlock libkbfs.DirBlock
			err = config.BlockOps().Get(
				ctx, irmd, rootPtr, &dirBlock)
			if err == nil {
				fmt.Printf("Got no error when getting root block %s; not doing anything\n", rootPtr)
				return nil
			}
			fmt.Printf("Got error %s when getting root block %s, so revision %d is broken. Making successor...\n",
				err, rootPtr, irmd.Revision())
		} else {
			// This happens in the wild, but only for folders used
			// for journal-related testing early on.
			fmt.Printf("Root block pointer is invalid, so revision %d is broken. Making successor...\n",
				irmd.Revision())
		}
	}

	rmdNext, err := irmd.MakeSuccessor(config.Codec(), irmd.MdID(), true)
	if err != nil {
		return err
	}

	// TODO: Add an option to scan for and use the last known good
	// root block.
	_, info, readyBlockData, err :=
		libkbfs.ResetRootBlock(ctx, config, uid, rmdNext)
	if err != nil {
		return err
	}

	fmt.Printf(
		"Will put an empty root block for tlfID=%s with blockInfo=%s and bufLen=%d\n",
		rmdNext.TlfID(), info, readyBlockData.GetEncodedSize())
	fmt.Print("Will put MD:\n")
	err = mdDumpOneReadOnly(ctx, config, rmdNext.ReadOnly())
	if err != nil {
		return err
	}

	if dryRun {
		fmt.Print("Dry-run set; not doing anything\n")
		return nil
	}

	fmt.Print("Are you sure you want to continue? [y/N]: ")
	response, err := bufio.NewReader(os.Stdin).ReadString('\n')
	if err != nil {
		return err
	}
	response = strings.ToLower(strings.TrimSpace(response))
	if response != "y" {
		fmt.Printf("Didn't confirm; not doing anything\n")
		return nil
	}

	fmt.Printf("Putting block %s...\n", info)

	err = libkbfs.PutBlockCheckQuota(
		ctx, config.BlockServer(), config.Reporter(),
		rmdNext.TlfID(), info.BlockPointer, readyBlockData,
		handle.GetCanonicalName())
	if err != nil {
		return err
	}

	// Assume there's no need to unembed the block changes.

	fmt.Printf("Putting revision %d...\n", rmdNext.Revision())

	mdID, err := config.MDOps().Put(ctx, rmdNext)
	if err != nil {
		return err
	}

	fmt.Printf("New MD has id %s\n", mdID)

	return nil
}

const mdResetUsageStr = `Usage:
  kbfstool md reset /keybase/[public|private]/user1,assertion2

`

func mdReset(ctx context.Context, config libkbfs.Config, args []string) (exitStatus int) {
	flags := flag.NewFlagSet("kbfs md reset", flag.ContinueOnError)
	checkValid := flags.Bool("c", true, "If set, don't do anything if the existing root block is valid")
	dryRun := flags.Bool("d", false, "Dry run: don't actually do anything.")
	err := flags.Parse(args)
	if err != nil {
		printError("md reset", err)
		return 1
	}

	inputs := flags.Args()
	if len(inputs) != 1 {
		fmt.Print(mdResetUsageStr)
		return 1
	}

	err = mdResetOne(ctx, config, inputs[0], *checkValid, *dryRun)
	if err != nil {
		printError("md reset", err)
		return 1
	}

	fmt.Print("\n")

	return 0
}
<commit_msg>Add force<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/keybase/kbfs/libkbfs"
	"golang.org/x/net/context"
)

func mdResetOne(
	ctx context.Context, config libkbfs.Config, tlfPath string,
	checkValid, dryRun, force bool) error {
	handle, err := parseTLFPath(ctx, config.KBPKI(), tlfPath)
	if err != nil {
		return err
	}

	// This function is loosely adapted from
	// folderBranchOps.initMDLocked.

	username, uid, err := config.KBPKI().GetCurrentUserInfo(ctx)
	if err != nil {
		return err
	}

	// Make sure we're a writer before doing anything else.
	if !handle.IsWriter(uid) {
		return libkbfs.NewWriteAccessError(
			handle, username, handle.GetCanonicalPath())
	}

	fmt.Printf("Looking for unmerged branch...\n")

	_, unmergedIRMD, err := config.MDOps().GetForHandle(
		ctx, handle, libkbfs.Unmerged)
	if err != nil {
		return err
	}
	if unmergedIRMD != (libkbfs.ImmutableRootMetadata{}) {
		return fmt.Errorf(
			"%s has unmerged data; try unstaging it first",
			tlfPath)
	}

	fmt.Printf("Getting latest metadata...\n")

	_, irmd, err := config.MDOps().GetForHandle(
		ctx, handle, libkbfs.Merged)
	if err != nil {
		return err
	}

	if irmd == (libkbfs.ImmutableRootMetadata{}) {
		fmt.Printf("No TLF found for %q\n", tlfPath)
		return nil
	}

	if checkValid {
		rootPtr := irmd.Data().Dir.BlockPointer
		if rootPtr.Ref().IsValid() {
			var dirBlock libkbfs.DirBlock
			err = config.BlockOps().Get(
				ctx, irmd, rootPtr, &dirBlock)
			if err == nil {
				fmt.Printf("Got no error when getting root block %s; not doing anything\n", rootPtr)
				return nil
			}
			fmt.Printf("Got error %s when getting root block %s, so revision %d is broken. Making successor...\n",
				err, rootPtr, irmd.Revision())
		} else {
			// This happens in the wild, but only for folders used
			// for journal-related testing early on.
			fmt.Printf("Root block pointer is invalid, so revision %d is broken. Making successor...\n",
				irmd.Revision())
		}
	}

	rmdNext, err := irmd.MakeSuccessor(config.Codec(), irmd.MdID(), true)
	if err != nil {
		return err
	}

	// TODO: Add an option to scan for and use the last known good
	// root block.
	_, info, readyBlockData, err :=
		libkbfs.ResetRootBlock(ctx, config, uid, rmdNext)
	if err != nil {
		return err
	}

	fmt.Printf(
		"Will put an empty root block for tlfID=%s with blockInfo=%s and bufLen=%d\n",
		rmdNext.TlfID(), info, readyBlockData.GetEncodedSize())
	fmt.Print("Will put MD:\n")
	err = mdDumpOneReadOnly(ctx, config, rmdNext.ReadOnly())
	if err != nil {
		return err
	}

	if dryRun {
		fmt.Print("Dry-run set; not doing anything\n")
		return nil
	}

	if !force {
		fmt.Print("Are you sure you want to continue? [y/N]: ")
		response, err := bufio.NewReader(os.Stdin).ReadString('\n')
		if err != nil {
			return err
		}
		response = strings.ToLower(strings.TrimSpace(response))
		if response != "y" {
			fmt.Printf("Didn't confirm; not doing anything\n")
			return nil
		}
	}

	fmt.Printf("Putting block %s...\n", info)

	err = libkbfs.PutBlockCheckQuota(
		ctx, config.BlockServer(), config.Reporter(),
		rmdNext.TlfID(), info.BlockPointer, readyBlockData,
		handle.GetCanonicalName())
	if err != nil {
		return err
	}

	// Assume there's no need to unembed the block changes.

	fmt.Printf("Putting revision %d...\n", rmdNext.Revision())

	mdID, err := config.MDOps().Put(ctx, rmdNext)
	if err != nil {
		return err
	}

	fmt.Printf("New MD has id %s\n", mdID)

	return nil
}

const mdResetUsageStr = `Usage:
  kbfstool md reset /keybase/[public|private]/user1,assertion2

`

func mdReset(ctx context.Context, config libkbfs.Config, args []string) (exitStatus int) {
	flags := flag.NewFlagSet("kbfs md reset", flag.ContinueOnError)
	checkValid := flags.Bool("c", true, "If set, don't do anything if the existing root block is valid")
	dryRun := flags.Bool("d", false, "Dry run: don't actually do anything.")
	force := flags.Bool("f", false, "If set, skip confirmation prompt.")
	err := flags.Parse(args)
	if err != nil {
		printError("md reset", err)
		return 1
	}

	inputs := flags.Args()
	if len(inputs) != 1 {
		fmt.Print(mdResetUsageStr)
		return 1
	}

	err = mdResetOne(ctx, config, inputs[0], *checkValid, *dryRun, *force)
	if err != nil {
		printError("md reset", err)
		return 1
	}

	fmt.Print("\n")

	return 0
}
<|endoftext|>


<commit_before>package kit

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/caarlos0/env"
	"github.com/imdario/mergo"
)

// Configuration is the structure of a configuration for an environment. This will
// get loaded into a theme client to dictate it's actions.
type Configuration struct {
	Environment  string        `yaml:"-" json:"-" env:"-"`
	Password     string        `yaml:"password,omitempty" json:"password,omitempty" env:"THEMEKIT_PASSWORD"`
	ThemeID      string        `yaml:"theme_id,omitempty" json:"theme_id,omitempty" env:"THEMEKIT_THEME_ID"`
	Domain       string        `yaml:"store" json:"store" env:"THEMEKIT_STORE"`
	Directory    string        `yaml:"directory,omitempty" json:"directory" env:"THEMEKIT_DIRECTORY"`
	IgnoredFiles []string      `yaml:"ignore_files,omitempty" json:"ignore_files,omitempty" env:"THEMEKIT_IGNORE_FILES" envSeparator:":"`
	Proxy        string        `yaml:"proxy,omitempty" json:"proxy,omitempty" env:"THEMEKIT_PROXY"`
	Ignores      []string      `yaml:"ignores,omitempty" json:"ignores,omitempty" env:"THEMEKIT_IGNORES" envSeparator:":"`
	Timeout      time.Duration `yaml:"timeout,omitempty" json:"timeout,omitempty" env:"THEMEKIT_TIMEOUT"`
}

// DefaultTimeout is the default timeout to kill any stalled processes.
const DefaultTimeout = 30 * time.Second

var (
	defaultConfig     = Configuration{}
	environmentConfig = Configuration{}
	flagConfig        = Configuration{}
)

func init() {
	pwd, _ := os.Getwd()

	defaultConfig = Configuration{
		Directory: pwd,
		Timeout:   DefaultTimeout,
	}

	env.Parse(&environmentConfig)
}

// SetFlagConfig will set the configuration that is set by your applications flags.
// Set the flag config before inializing any theme clients so that the loaded
// configurations will have the proper config precedence.
func SetFlagConfig(config Configuration) {
	flagConfig = config
}

// NewConfiguration will format a Configuration that combines the config from env variables,
// flags. Then it will validate that config. It will return the
// formatted configuration along with any validation errors. The config precedence
// is flags, environment variables, then the config file.
func NewConfiguration() (*Configuration, error) {
	return (&Configuration{}).compile()
}

func (conf *Configuration) compile() (*Configuration, error) {
	newConfig := &Configuration{}
	mergo.Merge(newConfig, &flagConfig)
	mergo.Merge(newConfig, &environmentConfig)
	mergo.Merge(newConfig, conf)
	mergo.Merge(newConfig, &defaultConfig)
	return newConfig, newConfig.Validate()
}

// Validate will check the configuration for any problems that will cause theme kit
// to function incorrectly.
func (conf Configuration) Validate() error {
	errors := []string{}

	if conf.ThemeID == "" {
		errors = append(errors, "missing theme_id")
	} else if !conf.IsLive() {
		if _, err := strconv.ParseInt(conf.ThemeID, 10, 64); err != nil {
			errors = append(errors, "invalid theme_id")
		}
	}

	if err := conf.validateNoThemeID(); err != nil {
		errors = append(errors, err.Error())
	}

	if len(errors) > 0 {
		return fmt.Errorf("Invalid configuration: %v", strings.Join(errors, ","))
	}
	return nil
}

func (conf Configuration) validateNoThemeID() error {
	errors := []string{}

	if len(conf.Domain) == 0 {
		errors = append(errors, "missing store domain")
	} else if !strings.HasSuffix(conf.Domain, "myshopify.com") &&
		!strings.HasSuffix(conf.Domain, "myshopify.io") &&
		!strings.HasPrefix(conf.Domain, "http://127.0.0.1:") { // for testing
		errors = append(errors, "invalid store domain must end in '.myshopify.com'")
	}

	if len(conf.Password) == 0 {
		errors = append(errors, "missing password")
	}

	if len(errors) > 0 {
		return fmt.Errorf(strings.Join(errors, ","))
	}
	return nil
}

// IsLive will return true if the configurations theme id is set to live
func (conf Configuration) IsLive() bool {
	return strings.ToLower(strings.TrimSpace(conf.ThemeID)) == "live"
}

// String will return a formatted string with the information about this configuration
func (conf Configuration) String() string {
	return fmt.Sprintf(`
Password     %v
ThemeID      %v
Domain       %v
Directory    %v
IgnoredFiles %v
Proxy        %v
Ignores      %v
Timeout      %v
	`,
		conf.Password,
		conf.ThemeID,
		conf.Domain,
		conf.Directory,
		conf.IgnoredFiles,
		conf.Proxy,
		conf.Ignores,
		conf.Timeout)
}

func (conf Configuration) asYAML() *Configuration {
	if conf.Directory == defaultConfig.Directory {
		conf.Directory = ""
	}
	if conf.Timeout == defaultConfig.Timeout {
		conf.Timeout = 0
	}
	return &conf
}
<commit_msg>update json tags on configuration<commit_after>package kit

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/caarlos0/env"
	"github.com/imdario/mergo"
)

// Configuration is the structure of a configuration for an environment. This will
// get loaded into a theme client to dictate it's actions.
type Configuration struct {
	Environment  string        `yaml:"-" json:"-" env:"-"`
	Password     string        `yaml:"password,omitempty" json:"password,omitempty" env:"THEMEKIT_PASSWORD"`
	ThemeID      string        `yaml:"theme_id,omitempty" json:"theme_id,omitempty" env:"THEMEKIT_THEME_ID"`
	Domain       string        `yaml:"store" json:"store" env:"THEMEKIT_STORE"`
	Directory    string        `yaml:"directory,omitempty" json:"directory,omitempty" env:"THEMEKIT_DIRECTORY"`
	IgnoredFiles []string      `yaml:"ignore_files,omitempty" json:"ignore_files,omitempty" env:"THEMEKIT_IGNORE_FILES" envSeparator:":"`
	Proxy        string        `yaml:"proxy,omitempty" json:"proxy,omitempty" env:"THEMEKIT_PROXY"`
	Ignores      []string      `yaml:"ignores,omitempty" json:"ignores,omitempty" env:"THEMEKIT_IGNORES" envSeparator:":"`
	Timeout      time.Duration `yaml:"timeout,omitempty" json:"timeout,omitempty" env:"THEMEKIT_TIMEOUT"`
}

// DefaultTimeout is the default timeout to kill any stalled processes.
const DefaultTimeout = 30 * time.Second

var (
	defaultConfig     = Configuration{}
	environmentConfig = Configuration{}
	flagConfig        = Configuration{}
)

func init() {
	pwd, _ := os.Getwd()

	defaultConfig = Configuration{
		Directory: pwd,
		Timeout:   DefaultTimeout,
	}

	env.Parse(&environmentConfig)
}

// SetFlagConfig will set the configuration that is set by your applications flags.
// Set the flag config before inializing any theme clients so that the loaded
// configurations will have the proper config precedence.
func SetFlagConfig(config Configuration) {
	flagConfig = config
}

// NewConfiguration will format a Configuration that combines the config from env variables,
// flags. Then it will validate that config. It will return the
// formatted configuration along with any validation errors. The config precedence
// is flags, environment variables, then the config file.
func NewConfiguration() (*Configuration, error) {
	return (&Configuration{}).compile()
}

func (conf *Configuration) compile() (*Configuration, error) {
	newConfig := &Configuration{}
	mergo.Merge(newConfig, &flagConfig)
	mergo.Merge(newConfig, &environmentConfig)
	mergo.Merge(newConfig, conf)
	mergo.Merge(newConfig, &defaultConfig)
	return newConfig, newConfig.Validate()
}

// Validate will check the configuration for any problems that will cause theme kit
// to function incorrectly.
func (conf Configuration) Validate() error {
	errors := []string{}

	if conf.ThemeID == "" {
		errors = append(errors, "missing theme_id")
	} else if !conf.IsLive() {
		if _, err := strconv.ParseInt(conf.ThemeID, 10, 64); err != nil {
			errors = append(errors, "invalid theme_id")
		}
	}

	if err := conf.validateNoThemeID(); err != nil {
		errors = append(errors, err.Error())
	}

	if len(errors) > 0 {
		return fmt.Errorf("Invalid configuration: %v", strings.Join(errors, ","))
	}
	return nil
}

func (conf Configuration) validateNoThemeID() error {
	errors := []string{}

	if len(conf.Domain) == 0 {
		errors = append(errors, "missing store domain")
	} else if !strings.HasSuffix(conf.Domain, "myshopify.com") &&
		!strings.HasSuffix(conf.Domain, "myshopify.io") &&
		!strings.HasPrefix(conf.Domain, "http://127.0.0.1:") { // for testing
		errors = append(errors, "invalid store domain must end in '.myshopify.com'")
	}

	if len(conf.Password) == 0 {
		errors = append(errors, "missing password")
	}

	if len(errors) > 0 {
		return fmt.Errorf(strings.Join(errors, ","))
	}
	return nil
}

// IsLive will return true if the configurations theme id is set to live
func (conf Configuration) IsLive() bool {
	return strings.ToLower(strings.TrimSpace(conf.ThemeID)) == "live"
}

// String will return a formatted string with the information about this configuration
func (conf Configuration) String() string {
	return fmt.Sprintf(`
Password     %v
ThemeID      %v
Domain       %v
Directory    %v
IgnoredFiles %v
Proxy        %v
Ignores      %v
Timeout      %v
	`,
		conf.Password,
		conf.ThemeID,
		conf.Domain,
		conf.Directory,
		conf.IgnoredFiles,
		conf.Proxy,
		conf.Ignores,
		conf.Timeout)
}

func (conf Configuration) asYAML() *Configuration {
	if conf.Directory == defaultConfig.Directory {
		conf.Directory = ""
	}
	if conf.Timeout == defaultConfig.Timeout {
		conf.Timeout = 0
	}
	return &conf
}
<|endoftext|>


<commit_before>package aero

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/OneOfOne/xxhash"
	"github.com/aerogo/session"
	"github.com/fatih/color"
	"github.com/julienschmidt/httprouter"
	cache "github.com/patrickmn/go-cache"
	"github.com/tomasen/realip"
	"github.com/valyala/fasthttp"
)

// This should be close to the MTU size of a TCP packet.
// Regarding performance it makes no sense to compress smaller files.
// Bandwidth can be saved however the savings are minimal for small files
// and the overhead of compressing can lead up to a 75% reduction
// in server speed under high load. Therefore in this case
// we're trying to optimize for performance, not bandwidth.
const gzipThreshold = 1450

const (
	serverHeader                  = "Server"
	server                        = "Aero"
	cacheControlHeader            = "Cache-Control"
	cacheControlAlwaysValidate    = "must-revalidate"
	cacheControlMedia             = "public, max-age=864000"
	contentTypeOptionsHeader      = "X-Content-Type-Options"
	contentTypeOptions            = "nosniff"
	xssProtectionHeader           = "X-XSS-Protection"
	xssProtection                 = "1; mode=block"
	etagHeader                    = "ETag"
	contentTypeHeader             = "Content-Type"
	contentTypeHTML               = "text/html; charset=utf-8"
	contentTypeJSON               = "application/json; charset=utf-8"
	contentTypePlainText          = "text/plain; charset=utf-8"
	contentEncodingHeader         = "Content-Encoding"
	contentEncodingGzip           = "gzip"
	contentLengthHeader           = "Content-Length"
	responseTimeHeader            = "X-Response-Time"
	ifNoneMatchHeader             = "If-None-Match"
	xFrameOptionsHeader           = "X-Frame-Options"
	xFrameOptions                 = "SAMEORIGIN"
	referrerPolicyHeader          = "Referrer-Policy"
	referrerPolicySameOrigin      = "no-referrer"
	strictTransportSecurityHeader = "Strict-Transport-Security"
	strictTransportSecurity       = "max-age=31536000; includeSubDomains; preload"
	contentSecurityPolicyHeader   = "Content-Security-Policy"
)

// Context ...
type Context struct {
	// net/http
	request  *http.Request
	response http.ResponseWriter
	params   httprouter.Params

	// A pointer to the application this request occured on.
	App *Application

	// Status code
	StatusCode int

	// Start time
	start time.Time

	// User session
	session *session.Session
}

// Handle ...
type Handle func(*Context) string

// Session returns the session of the context or creates and caches a new session.
func (ctx *Context) Session() *session.Session {
	// Return cached session if available.
	if ctx.session != nil {
		return ctx.session
	}

	// Check if the client has a session cookie already.
	cookie, err := ctx.request.Cookie("sid")

	if err == nil {
		sid := cookie.Value

		if session.IsValidID(sid) {
			ctx.session, err = ctx.App.Sessions.Store.Get(sid)

			if err != nil {
				color.Red(err.Error())
			}

			if ctx.session != nil {
				return ctx.session
			}
		}
	}

	// Create a new session
	ctx.session = ctx.App.Sessions.New()

	// Create a session cookie in the client
	ctx.createSessionCookie()

	return ctx.session
}

// createSessionCookie creates a session cookie in the client.
func (ctx *Context) createSessionCookie() {
	sessionCookie := http.Cookie{
		Name:     "sid",
		Value:    ctx.session.ID(),
		HttpOnly: true,
		Secure:   true,
		MaxAge:   ctx.App.Sessions.Duration,
		Path:     "/",
	}

	http.SetCookie(ctx.response, &sessionCookie)

	// HACK: Add SameSite attribute
	// Remove this once it's available inside http.Cookie
	cookieData := ctx.response.Header().Get("Set-Cookie")
	cookieData += "; SameSite"
	ctx.response.Header().Set("Set-Cookie", cookieData)
}

// HasSession indicates whether the client has a valid session or not.
func (ctx *Context) HasSession() bool {
	if ctx.session != nil {
		return true
	}

	cookie, err := ctx.request.Cookie("sid")

	if err != nil || !session.IsValidID(cookie.Value) {
		return false
	}

	ctx.session, _ = ctx.App.Sessions.Store.Get(cookie.Value)

	return ctx.session != nil
}

// JSON encodes the object to a JSON string and responds.
func (ctx *Context) JSON(value interface{}) string {
	bytes, _ := json.Marshal(value)

	ctx.SetResponseHeader(contentTypeHeader, contentTypeJSON)
	return string(bytes)
}

// HTML sends a HTML string.
func (ctx *Context) HTML(html string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	ctx.SetResponseHeader(contentTypeOptionsHeader, contentTypeOptions)
	ctx.SetResponseHeader(xssProtectionHeader, xssProtection)
	ctx.SetResponseHeader(xFrameOptionsHeader, xFrameOptions)
	ctx.SetResponseHeader(referrerPolicyHeader, referrerPolicySameOrigin)

	if ctx.App.Security.Certificate != "" {
		ctx.SetResponseHeader(strictTransportSecurityHeader, strictTransportSecurity)
		ctx.SetResponseHeader(contentSecurityPolicyHeader, ctx.App.contentSecurityPolicy)
	}

	return html
}

// Text sends a plain text string.
func (ctx *Context) Text(text string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypePlainText)
	return text
}

// File sends the contents of a local file and determines its mime type by extension.
func (ctx *Context) File(file string) string {
	extension := filepath.Ext(file)
	mimeType := mime.TypeByExtension(extension)
	data, _ := ioutil.ReadFile(file)

	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}

	ctx.SetResponseHeader(contentTypeHeader, mimeType)
	return string(data)
}

// TryWebP tries to serve a WebP image but will fall back to the specified extension if needed.
func (ctx *Context) TryWebP(path string, extension string) string {
	if ctx.CanUseWebP() {
		extension = ".webp"
	}

	return ctx.File(path + extension)
}

// Error should be used for sending error messages to the user.
func (ctx *Context) Error(statusCode int, explanation string, err error) string {
	ctx.StatusCode = statusCode
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	// ctx.App.Logger.Error(
	// 	color.RedString(explanation),
	// 	zap.String("error", err.Error()),
	// 	zap.String("url", ctx.request.RequestURI),
	// )
	detailed := err.Error()
	color.Red(detailed)
	return fmt.Sprintf("%s (%s)", explanation, detailed)
}

// GetRequestHeader retrieves the value for the request header.
func (ctx *Context) GetRequestHeader(header string) string {
	return ctx.request.Header.Get(header)
}

// SetResponseHeader sets response header to value.
func (ctx *Context) SetResponseHeader(header string, value string) {
	ctx.response.Header().Set(header, value)
}

// URI returns the relative path, e.g. /blog/post/123.
func (ctx *Context) URI() string {
	return ctx.request.URL.Path
}

// SetURI sets the relative path, e.g. /blog/post/123.
func (ctx *Context) SetURI(b string) {
	ctx.request.URL.Path = b
}

// Get retrieves an URL parameter.
func (ctx *Context) Get(param string) string {
	return ctx.params.ByName(param)
}

// GetInt retrieves an URL parameter as an integer.
func (ctx *Context) GetInt(param string) (int, error) {
	return strconv.Atoi(ctx.Get(param))
}

// RequestBody returns the request body as a string.
func (ctx *Context) RequestBody() []byte {
	body, err := ioutil.ReadAll(ctx.request.Body)

	if err != nil {
		panic(err)
	}

	return body
}

// RealIP tries to determine the real IP address of the request.
func (ctx *Context) RealIP() string {
	return realip.RealIP(ctx.request)
}

// UserAgent retrieves the user agent for the given request.
func (ctx *Context) UserAgent() string {
	ctx.request.URL.Query()
	return ctx.request.UserAgent()
}

// Query retrieves the value for the given URL query parameter.
func (ctx *Context) Query(param string) string {
	return ctx.request.URL.Query().Get(param)
}

// Redirect redirects to the given URL using status code 302.
func (ctx *Context) Redirect(url string) string {
	ctx.StatusCode = http.StatusFound
	ctx.SetResponseHeader("Location", url)
	return ""
}

// RedirectPermanently redirects to the given URL and indicates that this is a permanent change using status code 301.
func (ctx *Context) RedirectPermanently(url string) string {
	ctx.StatusCode = http.StatusPermanentRedirect
	ctx.SetResponseHeader("Location", url)
	return ""
}

// CanUseWebP checks the Accept header to find out if WebP is supported by the client's browser.
func (ctx *Context) CanUseWebP() bool {
	accept := ctx.GetRequestHeader("Accept")

	if strings.Index(accept, "image/webp") != -1 {
		return true
	}

	return false
}

// Respond responds either with raw code or gzipped if the
// code length is greater than the gzip threshold.
func (ctx *Context) Respond(code string) {
	ctx.RespondBytes(StringToBytesUnsafe(code))
}

// RespondBytes responds either with raw code or gzipped if the
// code length is greater than the gzip threshold. Requires a byte slice.
func (ctx *Context) RespondBytes(b []byte) {
	response := ctx.response
	header := response.Header()
	contentType := ctx.response.Header().Get(contentTypeHeader)
	isMedia := strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "video/")

	// Headers
	if isMedia {
		header.Set(cacheControlHeader, cacheControlMedia)
	} else {
		header.Set(cacheControlHeader, cacheControlAlwaysValidate)
		header.Set(serverHeader, server)
		header.Set(responseTimeHeader, strconv.FormatInt(time.Since(ctx.start).Nanoseconds()/1000, 10)+" us")
	}

	// Small response
	if len(b) < gzipThreshold {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// ETag generation
	h := xxhash.NewS64(0)
	h.Write(b)
	etag := strconv.FormatUint(h.Sum64(), 16)

	// If client cache is up to date, send 304 with no response body.
	clientETag := ctx.request.Header.Get(ifNoneMatchHeader)

	if etag == clientETag {
		response.WriteHeader(304)
		return
	}

	// Set ETag
	header.Set(etagHeader, etag)

	// No GZip?
	if !ctx.App.Config.GZip || isMedia {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// GZip
	header.Set(contentEncodingHeader, contentEncodingGzip)

	if ctx.App.Config.GZipCache {
		cachedResponse, found := ctx.App.gzipCache.Get(etag)

		if found {
			cachedResponseBytes := cachedResponse.([]byte)
			header.Set(contentLengthHeader, strconv.Itoa(len(cachedResponseBytes)))
			response.WriteHeader(ctx.StatusCode)
			response.Write(cachedResponseBytes)
			return
		}
	}

	var buffer bytes.Buffer
	writer := bufio.NewWriter(&buffer)
	fasthttp.WriteGzipLevel(writer, b, 9)
	writer.Flush()
	gzippedBytes := buffer.Bytes()

	header.Set(contentLengthHeader, strconv.Itoa(len(gzippedBytes)))
	response.WriteHeader(ctx.StatusCode)
	response.Write(gzippedBytes)

	if ctx.App.Config.GZipCache {
		ctx.App.gzipCache.Set(etag, gzippedBytes, cache.DefaultExpiration)
	}
}
<commit_msg>Specific lax value<commit_after>package aero

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/OneOfOne/xxhash"
	"github.com/aerogo/session"
	"github.com/fatih/color"
	"github.com/julienschmidt/httprouter"
	cache "github.com/patrickmn/go-cache"
	"github.com/tomasen/realip"
	"github.com/valyala/fasthttp"
)

// This should be close to the MTU size of a TCP packet.
// Regarding performance it makes no sense to compress smaller files.
// Bandwidth can be saved however the savings are minimal for small files
// and the overhead of compressing can lead up to a 75% reduction
// in server speed under high load. Therefore in this case
// we're trying to optimize for performance, not bandwidth.
const gzipThreshold = 1450

const (
	serverHeader                  = "Server"
	server                        = "Aero"
	cacheControlHeader            = "Cache-Control"
	cacheControlAlwaysValidate    = "must-revalidate"
	cacheControlMedia             = "public, max-age=864000"
	contentTypeOptionsHeader      = "X-Content-Type-Options"
	contentTypeOptions            = "nosniff"
	xssProtectionHeader           = "X-XSS-Protection"
	xssProtection                 = "1; mode=block"
	etagHeader                    = "ETag"
	contentTypeHeader             = "Content-Type"
	contentTypeHTML               = "text/html; charset=utf-8"
	contentTypeJSON               = "application/json; charset=utf-8"
	contentTypePlainText          = "text/plain; charset=utf-8"
	contentEncodingHeader         = "Content-Encoding"
	contentEncodingGzip           = "gzip"
	contentLengthHeader           = "Content-Length"
	responseTimeHeader            = "X-Response-Time"
	ifNoneMatchHeader             = "If-None-Match"
	xFrameOptionsHeader           = "X-Frame-Options"
	xFrameOptions                 = "SAMEORIGIN"
	referrerPolicyHeader          = "Referrer-Policy"
	referrerPolicySameOrigin      = "no-referrer"
	strictTransportSecurityHeader = "Strict-Transport-Security"
	strictTransportSecurity       = "max-age=31536000; includeSubDomains; preload"
	contentSecurityPolicyHeader   = "Content-Security-Policy"
)

// Context ...
type Context struct {
	// net/http
	request  *http.Request
	response http.ResponseWriter
	params   httprouter.Params

	// A pointer to the application this request occured on.
	App *Application

	// Status code
	StatusCode int

	// Start time
	start time.Time

	// User session
	session *session.Session
}

// Handle ...
type Handle func(*Context) string

// Session returns the session of the context or creates and caches a new session.
func (ctx *Context) Session() *session.Session {
	// Return cached session if available.
	if ctx.session != nil {
		return ctx.session
	}

	// Check if the client has a session cookie already.
	cookie, err := ctx.request.Cookie("sid")

	if err == nil {
		sid := cookie.Value

		if session.IsValidID(sid) {
			ctx.session, err = ctx.App.Sessions.Store.Get(sid)

			if err != nil {
				color.Red(err.Error())
			}

			if ctx.session != nil {
				return ctx.session
			}
		}
	}

	// Create a new session
	ctx.session = ctx.App.Sessions.New()

	// Create a session cookie in the client
	ctx.createSessionCookie()

	return ctx.session
}

// createSessionCookie creates a session cookie in the client.
func (ctx *Context) createSessionCookie() {
	sessionCookie := http.Cookie{
		Name:     "sid",
		Value:    ctx.session.ID(),
		HttpOnly: true,
		Secure:   true,
		MaxAge:   ctx.App.Sessions.Duration,
		Path:     "/",
	}

	http.SetCookie(ctx.response, &sessionCookie)

	// HACK: Add SameSite attribute
	// Remove this once it's available inside http.Cookie
	cookieData := ctx.response.Header().Get("Set-Cookie")
	cookieData += "; SameSite=lax"
	ctx.response.Header().Set("Set-Cookie", cookieData)
}

// HasSession indicates whether the client has a valid session or not.
func (ctx *Context) HasSession() bool {
	if ctx.session != nil {
		return true
	}

	cookie, err := ctx.request.Cookie("sid")

	if err != nil || !session.IsValidID(cookie.Value) {
		return false
	}

	ctx.session, _ = ctx.App.Sessions.Store.Get(cookie.Value)

	return ctx.session != nil
}

// JSON encodes the object to a JSON string and responds.
func (ctx *Context) JSON(value interface{}) string {
	bytes, _ := json.Marshal(value)

	ctx.SetResponseHeader(contentTypeHeader, contentTypeJSON)
	return string(bytes)
}

// HTML sends a HTML string.
func (ctx *Context) HTML(html string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	ctx.SetResponseHeader(contentTypeOptionsHeader, contentTypeOptions)
	ctx.SetResponseHeader(xssProtectionHeader, xssProtection)
	ctx.SetResponseHeader(xFrameOptionsHeader, xFrameOptions)
	ctx.SetResponseHeader(referrerPolicyHeader, referrerPolicySameOrigin)

	if ctx.App.Security.Certificate != "" {
		ctx.SetResponseHeader(strictTransportSecurityHeader, strictTransportSecurity)
		ctx.SetResponseHeader(contentSecurityPolicyHeader, ctx.App.contentSecurityPolicy)
	}

	return html
}

// Text sends a plain text string.
func (ctx *Context) Text(text string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypePlainText)
	return text
}

// File sends the contents of a local file and determines its mime type by extension.
func (ctx *Context) File(file string) string {
	extension := filepath.Ext(file)
	mimeType := mime.TypeByExtension(extension)
	data, _ := ioutil.ReadFile(file)

	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}

	ctx.SetResponseHeader(contentTypeHeader, mimeType)
	return string(data)
}

// TryWebP tries to serve a WebP image but will fall back to the specified extension if needed.
func (ctx *Context) TryWebP(path string, extension string) string {
	if ctx.CanUseWebP() {
		extension = ".webp"
	}

	return ctx.File(path + extension)
}

// Error should be used for sending error messages to the user.
func (ctx *Context) Error(statusCode int, explanation string, err error) string {
	ctx.StatusCode = statusCode
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	// ctx.App.Logger.Error(
	// 	color.RedString(explanation),
	// 	zap.String("error", err.Error()),
	// 	zap.String("url", ctx.request.RequestURI),
	// )
	detailed := err.Error()
	color.Red(detailed)
	return fmt.Sprintf("%s (%s)", explanation, detailed)
}

// GetRequestHeader retrieves the value for the request header.
func (ctx *Context) GetRequestHeader(header string) string {
	return ctx.request.Header.Get(header)
}

// SetResponseHeader sets response header to value.
func (ctx *Context) SetResponseHeader(header string, value string) {
	ctx.response.Header().Set(header, value)
}

// URI returns the relative path, e.g. /blog/post/123.
func (ctx *Context) URI() string {
	return ctx.request.URL.Path
}

// SetURI sets the relative path, e.g. /blog/post/123.
func (ctx *Context) SetURI(b string) {
	ctx.request.URL.Path = b
}

// Get retrieves an URL parameter.
func (ctx *Context) Get(param string) string {
	return ctx.params.ByName(param)
}

// GetInt retrieves an URL parameter as an integer.
func (ctx *Context) GetInt(param string) (int, error) {
	return strconv.Atoi(ctx.Get(param))
}

// RequestBody returns the request body as a string.
func (ctx *Context) RequestBody() []byte {
	body, err := ioutil.ReadAll(ctx.request.Body)

	if err != nil {
		panic(err)
	}

	return body
}

// RealIP tries to determine the real IP address of the request.
func (ctx *Context) RealIP() string {
	return realip.RealIP(ctx.request)
}

// UserAgent retrieves the user agent for the given request.
func (ctx *Context) UserAgent() string {
	ctx.request.URL.Query()
	return ctx.request.UserAgent()
}

// Query retrieves the value for the given URL query parameter.
func (ctx *Context) Query(param string) string {
	return ctx.request.URL.Query().Get(param)
}

// Redirect redirects to the given URL using status code 302.
func (ctx *Context) Redirect(url string) string {
	ctx.StatusCode = http.StatusFound
	ctx.SetResponseHeader("Location", url)
	return ""
}

// RedirectPermanently redirects to the given URL and indicates that this is a permanent change using status code 301.
func (ctx *Context) RedirectPermanently(url string) string {
	ctx.StatusCode = http.StatusPermanentRedirect
	ctx.SetResponseHeader("Location", url)
	return ""
}

// CanUseWebP checks the Accept header to find out if WebP is supported by the client's browser.
func (ctx *Context) CanUseWebP() bool {
	accept := ctx.GetRequestHeader("Accept")

	if strings.Index(accept, "image/webp") != -1 {
		return true
	}

	return false
}

// Respond responds either with raw code or gzipped if the
// code length is greater than the gzip threshold.
func (ctx *Context) Respond(code string) {
	ctx.RespondBytes(StringToBytesUnsafe(code))
}

// RespondBytes responds either with raw code or gzipped if the
// code length is greater than the gzip threshold. Requires a byte slice.
func (ctx *Context) RespondBytes(b []byte) {
	response := ctx.response
	header := response.Header()
	contentType := ctx.response.Header().Get(contentTypeHeader)
	isMedia := strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "video/")

	// Headers
	if isMedia {
		header.Set(cacheControlHeader, cacheControlMedia)
	} else {
		header.Set(cacheControlHeader, cacheControlAlwaysValidate)
		header.Set(serverHeader, server)
		header.Set(responseTimeHeader, strconv.FormatInt(time.Since(ctx.start).Nanoseconds()/1000, 10)+" us")
	}

	// Small response
	if len(b) < gzipThreshold {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// ETag generation
	h := xxhash.NewS64(0)
	h.Write(b)
	etag := strconv.FormatUint(h.Sum64(), 16)

	// If client cache is up to date, send 304 with no response body.
	clientETag := ctx.request.Header.Get(ifNoneMatchHeader)

	if etag == clientETag {
		response.WriteHeader(304)
		return
	}

	// Set ETag
	header.Set(etagHeader, etag)

	// No GZip?
	if !ctx.App.Config.GZip || isMedia {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// GZip
	header.Set(contentEncodingHeader, contentEncodingGzip)

	if ctx.App.Config.GZipCache {
		cachedResponse, found := ctx.App.gzipCache.Get(etag)

		if found {
			cachedResponseBytes := cachedResponse.([]byte)
			header.Set(contentLengthHeader, strconv.Itoa(len(cachedResponseBytes)))
			response.WriteHeader(ctx.StatusCode)
			response.Write(cachedResponseBytes)
			return
		}
	}

	var buffer bytes.Buffer
	writer := bufio.NewWriter(&buffer)
	fasthttp.WriteGzipLevel(writer, b, 9)
	writer.Flush()
	gzippedBytes := buffer.Bytes()

	header.Set(contentLengthHeader, strconv.Itoa(len(gzippedBytes)))
	response.WriteHeader(ctx.StatusCode)
	response.Write(gzippedBytes)

	if ctx.App.Config.GZipCache {
		ctx.App.gzipCache.Set(etag, gzippedBytes, cache.DefaultExpiration)
	}
}
<|endoftext|>