<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2019 Red Hat, Inc.
 *
 */

package webhooks

import (
	"crypto/x509"
	"fmt"
	"sync"

	k8sv1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/cert"

	"kubevirt.io/kubevirt/pkg/virt-operator/creation/components"

	"kubevirt.io/kubevirt/pkg/util"
)

type ClientCAManager interface {
	GetCurrent() (*x509.CertPool, error)
}

type manager struct {
	store        cache.Store
	lock         *sync.Mutex
	lastRevision string
	namespace    string
	name         string
	secretKey    string

	lastPool *x509.CertPool
}

func NewKubernetesClientCAManager(configMapCache cache.Store) ClientCAManager {
	return &manager{
		store:     configMapCache,
		lock:      &sync.Mutex{},
		namespace: metav1.NamespaceSystem,
		name:      util.ExtensionAPIServerAuthenticationConfigMap,
		secretKey: util.RequestHeaderClientCAFileKey,
	}
}

func NewCAManager(configMapCache cache.Store, namespace string, configMapName string) ClientCAManager {
	return &manager{
		store:     configMapCache,
		lock:      &sync.Mutex{},
		namespace: namespace,
		name:      configMapName,
		secretKey: components.CABundleKey,
	}
}

func (m *manager) GetCurrent() (*x509.CertPool, error) {
	m.lock.Lock()
	defer m.lock.Unlock()

	obj, exists, err := m.store.GetByKey(m.namespace + "/" + m.name)

	if err != nil {
		return nil, err
	} else if !exists {
		if m.lastPool != nil {
			return m.lastPool, nil
		}

		return nil, fmt.Errorf("configmap %s not found. Unable to detect request header CA", m.name)
	}

	configMap := obj.(*k8sv1.ConfigMap)

	// no change detected.
	if m.lastRevision == configMap.ResourceVersion {
		return m.lastPool, nil
	}

	requestHeaderClientCA, ok := configMap.Data[m.secretKey]
	if !ok {
		return nil, fmt.Errorf("requestheader-client-ca-file not found in extension-apiserver-authentication ConfigMap")
	}

	certs, err := cert.ParseCertsPEM([]byte(requestHeaderClientCA))
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	for _, cert := range certs {
		pool.AddCert(cert)
	}

	m.lastRevision = configMap.ResourceVersion
	m.lastPool = pool

	return pool, nil
}
<commit_msg>Add a log message if we pick up a new CA bundle<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2019 Red Hat, Inc.
 *
 */

package webhooks

import (
	"crypto/x509"
	"fmt"
	"sync"

	k8sv1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/cert"

	"kubevirt.io/client-go/log"
	"kubevirt.io/kubevirt/pkg/util"
	"kubevirt.io/kubevirt/pkg/virt-operator/creation/components"
)

type ClientCAManager interface {
	GetCurrent() (*x509.CertPool, error)
}

type manager struct {
	store        cache.Store
	lock         *sync.Mutex
	lastRevision string
	namespace    string
	name         string
	secretKey    string

	lastPool *x509.CertPool
}

func NewKubernetesClientCAManager(configMapCache cache.Store) ClientCAManager {
	return &manager{
		store:        configMapCache,
		lock:         &sync.Mutex{},
		namespace:    metav1.NamespaceSystem,
		name:         util.ExtensionAPIServerAuthenticationConfigMap,
		secretKey:    util.RequestHeaderClientCAFileKey,
		lastRevision: "-1",
	}
}

func NewCAManager(configMapCache cache.Store, namespace string, configMapName string) ClientCAManager {
	return &manager{
		store:        configMapCache,
		lock:         &sync.Mutex{},
		namespace:    namespace,
		name:         configMapName,
		secretKey:    components.CABundleKey,
		lastRevision: "-1",
	}
}

func (m *manager) GetCurrent() (*x509.CertPool, error) {
	m.lock.Lock()
	defer m.lock.Unlock()

	obj, exists, err := m.store.GetByKey(m.namespace + "/" + m.name)

	if err != nil {
		return nil, err
	} else if !exists {
		if m.lastPool != nil {
			return m.lastPool, nil
		}

		return nil, fmt.Errorf("configmap %s not found. Unable to detect request header CA", m.name)
	}

	configMap := obj.(*k8sv1.ConfigMap)

	// no change detected.
	if m.lastRevision == configMap.ResourceVersion {
		return m.lastPool, nil
	}
	log.DefaultLogger().Infof("CA update in configmap %s/%s detected. Updating from resource version %v to %v", m.namespace, m.name, m.lastRevision, configMap.ResourceVersion)

	requestHeaderClientCA, ok := configMap.Data[m.secretKey]
	if !ok {
		return nil, fmt.Errorf("requestheader-client-ca-file not found in configmap %s/%s", m.namespace, m.name)
	}

	certs, err := cert.ParseCertsPEM([]byte(requestHeaderClientCA))
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	for _, cert := range certs {
		pool.AddCert(cert)
	}

	m.lastRevision = configMap.ResourceVersion
	m.lastPool = pool

	return pool, nil
}
<|endoftext|>


<commit_before>package ssh

import (
	"fmt"
	"os/exec"

	. "github.com/onsi/gomega"

	. "github.com/onsi/ginkgo/v2"
)

var _ = Describe("Wrapped SSH", func() {

	var fakeKind, fakeNamespace, fakeName string
	var ssh SSH

	BeforeEach(func() {
		fakeKind = "fake-kind"
		fakeNamespace = "fake-ns"
		fakeName = "fake-name"
		ssh = SSH{}
	})

	Context("buildSSHTarget", func() {

		It("with SSH username", func() {
			ssh.options = SSHOptions{SshUsername: "testuser"}
			sshTarget := ssh.buildSSHTarget(fakeKind, fakeNamespace, fakeName)
			Expect(sshTarget).To(Equal("testuser@fake-kind/fake-name.fake-ns"))
		})

		It("without SSH username", func() {
			sshTarget := ssh.buildSSHTarget(fakeKind, fakeNamespace, fakeName)
			Expect(sshTarget).To(Equal("fake-kind/fake-name.fake-ns"))
		})

	})

	It("buildProxyCommandOption", func() {
		const sshPort = 12345
		ssh.options = SSHOptions{SshPort: sshPort}
		proxyCommand := ssh.buildProxyCommandOption(fakeKind, fakeNamespace, fakeName)
		expected := fmt.Sprintf("port-forward --stdio=true fake-kind/fake-name.fake-ns %d", sshPort)
		Expect(proxyCommand).To(ContainSubstring(expected))
	})

	It("runLocalCommandClient", func() {
		runCommand = func(cmd *exec.Cmd) error {
			Expect(cmd).ToNot(BeNil())
			Expect(cmd.Args).To(HaveLen(3))
			Expect(cmd.Args[0]).To(Equal("ssh"))
			Expect(cmd.Args[1]).To(Equal(ssh.buildProxyCommandOption(fakeKind, fakeNamespace, fakeName)))
			Expect(cmd.Args[2]).To(Equal(ssh.buildSSHTarget(fakeKind, fakeNamespace, fakeName)))

			return nil
		}

		err := ssh.runLocalCommandClient(fakeKind, fakeNamespace, fakeName)
		Expect(err).ShouldNot(HaveOccurred())
	})
})
<commit_msg>adjust unit tests for wrapped ssh<commit_after>package ssh

import (
	"fmt"
	"os/exec"

	. "github.com/onsi/gomega"

	. "github.com/onsi/ginkgo/v2"
)

var _ = Describe("Wrapped SSH", func() {

	var fakeKind, fakeNamespace, fakeName string
	var ssh SSH

	BeforeEach(func() {
		fakeKind = "fake-kind"
		fakeNamespace = "fake-ns"
		fakeName = "fake-name"
		ssh = SSH{}
	})

	Context("buildSSHTarget", func() {

		It("with SSH username", func() {
			ssh.options = SSHOptions{SshUsername: "testuser"}
			sshTarget := ssh.buildSSHTarget(fakeKind, fakeNamespace, fakeName)
			Expect(sshTarget[0]).To(Equal("testuser@fake-kind/fake-name.fake-ns"))
		})

		It("without SSH username", func() {
			sshTarget := ssh.buildSSHTarget(fakeKind, fakeNamespace, fakeName)
			Expect(sshTarget[0]).To(Equal("fake-kind/fake-name.fake-ns"))
		})

	})

	It("buildProxyCommandOption", func() {
		const sshPort = 12345
		ssh.options = SSHOptions{SshPort: sshPort}
		proxyCommand := ssh.buildProxyCommandOption(fakeKind, fakeNamespace, fakeName)
		expected := fmt.Sprintf("port-forward --stdio=true fake-kind/fake-name.fake-ns %d", sshPort)
		Expect(proxyCommand).To(ContainSubstring(expected))
	})

	It("runLocalCommandClient", func() {
		runCommand = func(cmd *exec.Cmd) error {
			Expect(cmd).ToNot(BeNil())
			Expect(cmd.Args).To(HaveLen(4))
			Expect(cmd.Args[0]).To(Equal("ssh"))
			Expect(cmd.Args[2]).To(Equal(ssh.buildProxyCommandOption(fakeKind, fakeNamespace, fakeName)))
			Expect(cmd.Args[3]).To(Equal(ssh.buildSSHTarget(fakeKind, fakeNamespace, fakeName)[0]))

			return nil
		}

		err := ssh.runLocalCommandClient(fakeKind, fakeNamespace, fakeName)
		Expect(err).ShouldNot(HaveOccurred())
	})
})
<|endoftext|>


<commit_before>//go:build linux
// +build linux

/*
Copyright 2020 The Kubernetes 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 volume

import (
	"fmt"
	"os"
	"path/filepath"
	"syscall"
	"testing"

	v1 "k8s.io/api/core/v1"
	utiltesting "k8s.io/client-go/util/testing"
)

type localFakeMounter struct {
	path       string
	attributes Attributes
}

func (l *localFakeMounter) GetPath() string {
	return l.path
}

func (l *localFakeMounter) GetAttributes() Attributes {
	return l.attributes
}

func (l *localFakeMounter) SetUp(mounterArgs MounterArgs) error {
	return nil
}

func (l *localFakeMounter) SetUpAt(dir string, mounterArgs MounterArgs) error {
	return nil
}

func (l *localFakeMounter) GetMetrics() (*Metrics, error) {
	return nil, nil
}

func TestSkipPermissionChange(t *testing.T) {
	always := v1.FSGroupChangeAlways
	onrootMismatch := v1.FSGroupChangeOnRootMismatch
	tests := []struct {
		description         string
		fsGroupChangePolicy *v1.PodFSGroupChangePolicy
		gidOwnerMatch       bool
		permissionMatch     bool
		sgidMatch           bool
		skipPermssion       bool
	}{
		{
			description:   "skippermission=false, policy=nil",
			skipPermssion: false,
		},
		{
			description:         "skippermission=false, policy=always",
			fsGroupChangePolicy: &always,
			skipPermssion:       false,
		},
		{
			description:         "skippermission=false, policy=always, gidmatch=true",
			fsGroupChangePolicy: &always,
			skipPermssion:       false,
			gidOwnerMatch:       true,
		},
		{
			description:         "skippermission=false, policy=nil, gidmatch=true",
			fsGroupChangePolicy: nil,
			skipPermssion:       false,
			gidOwnerMatch:       true,
		},
		{
			description:         "skippermission=false, policy=onrootmismatch, gidmatch=false",
			fsGroupChangePolicy: &onrootMismatch,
			gidOwnerMatch:       false,
			skipPermssion:       false,
		},
		{
			description:         "skippermission=false, policy=onrootmismatch, gidmatch=true, permmatch=false",
			fsGroupChangePolicy: &onrootMismatch,
			gidOwnerMatch:       true,
			permissionMatch:     false,
			skipPermssion:       false,
		},
		{
			description:         "skippermission=false, policy=onrootmismatch, gidmatch=true, permmatch=true",
			fsGroupChangePolicy: &onrootMismatch,
			gidOwnerMatch:       true,
			permissionMatch:     true,
			skipPermssion:       false,
		},
		{
			description:         "skippermission=false, policy=onrootmismatch, gidmatch=true, permmatch=true, sgidmatch=true",
			fsGroupChangePolicy: &onrootMismatch,
			gidOwnerMatch:       true,
			permissionMatch:     true,
			sgidMatch:           true,
			skipPermssion:       true,
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			tmpDir, err := utiltesting.MkTmpdir("volume_linux_test")
			if err != nil {
				t.Fatalf("error creating temp dir: %v", err)
			}

			defer os.RemoveAll(tmpDir)

			info, err := os.Lstat(tmpDir)
			if err != nil {
				t.Fatalf("error reading permission of tmpdir: %v", err)
			}

			stat, ok := info.Sys().(*syscall.Stat_t)
			if !ok || stat == nil {
				t.Fatalf("error reading permission stats for tmpdir: %s", tmpDir)
			}

			gid := stat.Gid

			var expectedGid int64

			if test.gidOwnerMatch {
				expectedGid = int64(gid)
			} else {
				expectedGid = int64(gid + 3000)
			}

			mask := rwMask

			if test.permissionMatch {
				mask |= execMask

			}
			if test.sgidMatch {
				mask |= os.ModeSetgid
				mask = info.Mode() | mask
			} else {
				nosgidPerm := info.Mode() &^ os.ModeSetgid
				mask = nosgidPerm | mask
			}

			err = os.Chmod(tmpDir, mask)
			if err != nil {
				t.Errorf("Chmod failed on %v: %v", tmpDir, err)
			}

			mounter := &localFakeMounter{path: tmpDir}
			ok = skipPermissionChange(mounter, &expectedGid, test.fsGroupChangePolicy)
			if ok != test.skipPermssion {
				t.Errorf("for %s expected skipPermission to be %v got %v", test.description, test.skipPermssion, ok)
			}

		})
	}
}

func TestSetVolumeOwnershipMode(t *testing.T) {
	always := v1.FSGroupChangeAlways
	onrootMismatch := v1.FSGroupChangeOnRootMismatch
	expectedMask := rwMask | os.ModeSetgid | execMask

	tests := []struct {
		description         string
		fsGroupChangePolicy *v1.PodFSGroupChangePolicy
		setupFunc           func(path string) error
		assertFunc          func(path string) error
	}{
		{
			description:         "featuregate=on, fsgroupchangepolicy=always",
			fsGroupChangePolicy: &always,
			setupFunc: func(path string) error {
				info, err := os.Lstat(path)
				if err != nil {
					return err
				}
				// change mode of root folder to be right
				err = os.Chmod(path, info.Mode()|expectedMask)
				if err != nil {
					return err
				}

				// create a subdirectory with invalid permissions
				rogueDir := filepath.Join(path, "roguedir")
				nosgidPerm := info.Mode() &^ os.ModeSetgid
				err = os.Mkdir(rogueDir, nosgidPerm)
				if err != nil {
					return err
				}
				return nil
			},
			assertFunc: func(path string) error {
				rogueDir := filepath.Join(path, "roguedir")
				hasCorrectPermissions := verifyDirectoryPermission(rogueDir, false /*readOnly*/)
				if !hasCorrectPermissions {
					return fmt.Errorf("invalid permissions on %s", rogueDir)
				}
				return nil
			},
		},
		{
			description:         "featuregate=on, fsgroupchangepolicy=onrootmismatch,rootdir=validperm",
			fsGroupChangePolicy: &onrootMismatch,
			setupFunc: func(path string) error {
				info, err := os.Lstat(path)
				if err != nil {
					return err
				}
				// change mode of root folder to be right
				err = os.Chmod(path, info.Mode()|expectedMask)
				if err != nil {
					return err
				}

				// create a subdirectory with invalid permissions
				rogueDir := filepath.Join(path, "roguedir")
				err = os.Mkdir(rogueDir, rwMask)
				if err != nil {
					return err
				}
				return nil
			},
			assertFunc: func(path string) error {
				rogueDir := filepath.Join(path, "roguedir")
				hasCorrectPermissions := verifyDirectoryPermission(rogueDir, false /*readOnly*/)
				if hasCorrectPermissions {
					return fmt.Errorf("invalid permissions on %s", rogueDir)
				}
				return nil
			},
		},
		{
			description:         "featuregate=on, fsgroupchangepolicy=onrootmismatch,rootdir=invalidperm",
			fsGroupChangePolicy: &onrootMismatch,
			setupFunc: func(path string) error {
				// change mode of root folder to be right
				err := os.Chmod(path, 0770)
				if err != nil {
					return err
				}

				// create a subdirectory with invalid permissions
				rogueDir := filepath.Join(path, "roguedir")
				err = os.Mkdir(rogueDir, rwMask)
				if err != nil {
					return err
				}
				return nil
			},
			assertFunc: func(path string) error {
				rogueDir := filepath.Join(path, "roguedir")
				hasCorrectPermissions := verifyDirectoryPermission(rogueDir, false /*readOnly*/)
				if !hasCorrectPermissions {
					return fmt.Errorf("invalid permissions on %s", rogueDir)
				}
				return nil
			},
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			tmpDir, err := utiltesting.MkTmpdir("volume_linux_ownership")
			if err != nil {
				t.Fatalf("error creating temp dir: %v", err)
			}

			defer os.RemoveAll(tmpDir)
			info, err := os.Lstat(tmpDir)
			if err != nil {
				t.Fatalf("error reading permission of tmpdir: %v", err)
			}

			stat, ok := info.Sys().(*syscall.Stat_t)
			if !ok || stat == nil {
				t.Fatalf("error reading permission stats for tmpdir: %s", tmpDir)
			}

			var expectedGid int64 = int64(stat.Gid)
			err = test.setupFunc(tmpDir)
			if err != nil {
				t.Errorf("for %s error running setup with: %v", test.description, err)
			}

			mounter := &localFakeMounter{path: tmpDir}
			err = SetVolumeOwnership(mounter, &expectedGid, test.fsGroupChangePolicy, nil)
			if err != nil {
				t.Errorf("for %s error changing ownership with: %v", test.description, err)
			}
			err = test.assertFunc(tmpDir)
			if err != nil {
				t.Errorf("for %s error verifying permissions with: %v", test.description, err)
			}
		})
	}
}

// verifyDirectoryPermission checks if given path has directory permissions
// that is expected by k8s. If returns true if it does otherwise false
func verifyDirectoryPermission(path string, readonly bool) bool {
	info, err := os.Lstat(path)
	if err != nil {
		return false
	}
	stat, ok := info.Sys().(*syscall.Stat_t)
	if !ok || stat == nil {
		return false
	}
	unixPerms := rwMask

	if readonly {
		unixPerms = roMask
	}

	unixPerms |= execMask
	filePerm := info.Mode().Perm()
	if (unixPerms&filePerm == unixPerms) && (info.Mode()&os.ModeSetgid != 0) {
		return true
	}
	return false
}

func TestSetVolumeOwnershipOwner(t *testing.T) {
	fsGroup := int64(3000)
	currentUid := os.Getuid()
	if currentUid != 0 {
		t.Skip("running as non-root")
	}
	currentGid := os.Getgid()

	tests := []struct {
		description string
		fsGroup     *int64
		setupFunc   func(path string) error
		assertFunc  func(path string) error
	}{
		{
			description: "fsGroup=nil",
			fsGroup:     nil,
			setupFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0755)
				if err != nil {
					return err
				}
				file.Close()
				return nil
			},
			assertFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				if !verifyFileOwner(filename, currentUid, currentGid) {
					return fmt.Errorf("invalid owner on %s", filename)
				}
				return nil
			},
		},
		{
			description: "*fsGroup=3000",
			fsGroup:     &fsGroup,
			setupFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0755)
				if err != nil {
					return err
				}
				file.Close()
				return nil
			},
			assertFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				if !verifyFileOwner(filename, currentUid, int(fsGroup)) {
					return fmt.Errorf("invalid owner on %s", filename)
				}
				return nil
			},
		},
		{
			description: "symlink",
			fsGroup:     &fsGroup,
			setupFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0755)
				if err != nil {
					return err
				}
				file.Close()

				symname := filepath.Join(path, "file_link.txt")
				err = os.Symlink(filename, symname)
				if err != nil {
					return err
				}

				return nil
			},
			assertFunc: func(path string) error {
				symname := filepath.Join(path, "file_link.txt")
				if !verifyFileOwner(symname, currentUid, int(fsGroup)) {
					return fmt.Errorf("invalid owner on %s", symname)
				}
				return nil
			},
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			tmpDir, err := utiltesting.MkTmpdir("volume_linux_ownership")
			if err != nil {
				t.Fatalf("error creating temp dir: %v", err)
			}

			defer os.RemoveAll(tmpDir)

			err = test.setupFunc(tmpDir)
			if err != nil {
				t.Errorf("for %s error running setup with: %v", test.description, err)
			}

			mounter := &localFakeMounter{path: tmpDir}
			always := v1.FSGroupChangeAlways
			err = SetVolumeOwnership(mounter, test.fsGroup, &always, nil)
			if err != nil {
				t.Errorf("for %s error changing ownership with: %v", test.description, err)
			}
			err = test.assertFunc(tmpDir)
			if err != nil {
				t.Errorf("for %s error verifying permissions with: %v", test.description, err)
			}
		})
	}
}

// verifyFileOwner checks if given path is owned by uid and gid.
// It returns true if it is otherwise false.
func verifyFileOwner(path string, uid, gid int) bool {
	info, err := os.Lstat(path)
	if err != nil {
		return false
	}
	stat, ok := info.Sys().(*syscall.Stat_t)
	if !ok || stat == nil {
		return false
	}

	if int(stat.Uid) != uid || int(stat.Gid) != gid {
		return false
	}

	return true
}
<commit_msg>volume: use the effective uid<commit_after>//go:build linux
// +build linux

/*
Copyright 2020 The Kubernetes 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 volume

import (
	"fmt"
	"os"
	"path/filepath"
	"syscall"
	"testing"

	v1 "k8s.io/api/core/v1"
	utiltesting "k8s.io/client-go/util/testing"
)

type localFakeMounter struct {
	path       string
	attributes Attributes
}

func (l *localFakeMounter) GetPath() string {
	return l.path
}

func (l *localFakeMounter) GetAttributes() Attributes {
	return l.attributes
}

func (l *localFakeMounter) SetUp(mounterArgs MounterArgs) error {
	return nil
}

func (l *localFakeMounter) SetUpAt(dir string, mounterArgs MounterArgs) error {
	return nil
}

func (l *localFakeMounter) GetMetrics() (*Metrics, error) {
	return nil, nil
}

func TestSkipPermissionChange(t *testing.T) {
	always := v1.FSGroupChangeAlways
	onrootMismatch := v1.FSGroupChangeOnRootMismatch
	tests := []struct {
		description         string
		fsGroupChangePolicy *v1.PodFSGroupChangePolicy
		gidOwnerMatch       bool
		permissionMatch     bool
		sgidMatch           bool
		skipPermssion       bool
	}{
		{
			description:   "skippermission=false, policy=nil",
			skipPermssion: false,
		},
		{
			description:         "skippermission=false, policy=always",
			fsGroupChangePolicy: &always,
			skipPermssion:       false,
		},
		{
			description:         "skippermission=false, policy=always, gidmatch=true",
			fsGroupChangePolicy: &always,
			skipPermssion:       false,
			gidOwnerMatch:       true,
		},
		{
			description:         "skippermission=false, policy=nil, gidmatch=true",
			fsGroupChangePolicy: nil,
			skipPermssion:       false,
			gidOwnerMatch:       true,
		},
		{
			description:         "skippermission=false, policy=onrootmismatch, gidmatch=false",
			fsGroupChangePolicy: &onrootMismatch,
			gidOwnerMatch:       false,
			skipPermssion:       false,
		},
		{
			description:         "skippermission=false, policy=onrootmismatch, gidmatch=true, permmatch=false",
			fsGroupChangePolicy: &onrootMismatch,
			gidOwnerMatch:       true,
			permissionMatch:     false,
			skipPermssion:       false,
		},
		{
			description:         "skippermission=false, policy=onrootmismatch, gidmatch=true, permmatch=true",
			fsGroupChangePolicy: &onrootMismatch,
			gidOwnerMatch:       true,
			permissionMatch:     true,
			skipPermssion:       false,
		},
		{
			description:         "skippermission=false, policy=onrootmismatch, gidmatch=true, permmatch=true, sgidmatch=true",
			fsGroupChangePolicy: &onrootMismatch,
			gidOwnerMatch:       true,
			permissionMatch:     true,
			sgidMatch:           true,
			skipPermssion:       true,
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			tmpDir, err := utiltesting.MkTmpdir("volume_linux_test")
			if err != nil {
				t.Fatalf("error creating temp dir: %v", err)
			}

			defer os.RemoveAll(tmpDir)

			info, err := os.Lstat(tmpDir)
			if err != nil {
				t.Fatalf("error reading permission of tmpdir: %v", err)
			}

			stat, ok := info.Sys().(*syscall.Stat_t)
			if !ok || stat == nil {
				t.Fatalf("error reading permission stats for tmpdir: %s", tmpDir)
			}

			gid := stat.Gid

			var expectedGid int64

			if test.gidOwnerMatch {
				expectedGid = int64(gid)
			} else {
				expectedGid = int64(gid + 3000)
			}

			mask := rwMask

			if test.permissionMatch {
				mask |= execMask

			}
			if test.sgidMatch {
				mask |= os.ModeSetgid
				mask = info.Mode() | mask
			} else {
				nosgidPerm := info.Mode() &^ os.ModeSetgid
				mask = nosgidPerm | mask
			}

			err = os.Chmod(tmpDir, mask)
			if err != nil {
				t.Errorf("Chmod failed on %v: %v", tmpDir, err)
			}

			mounter := &localFakeMounter{path: tmpDir}
			ok = skipPermissionChange(mounter, &expectedGid, test.fsGroupChangePolicy)
			if ok != test.skipPermssion {
				t.Errorf("for %s expected skipPermission to be %v got %v", test.description, test.skipPermssion, ok)
			}

		})
	}
}

func TestSetVolumeOwnershipMode(t *testing.T) {
	always := v1.FSGroupChangeAlways
	onrootMismatch := v1.FSGroupChangeOnRootMismatch
	expectedMask := rwMask | os.ModeSetgid | execMask

	tests := []struct {
		description         string
		fsGroupChangePolicy *v1.PodFSGroupChangePolicy
		setupFunc           func(path string) error
		assertFunc          func(path string) error
	}{
		{
			description:         "featuregate=on, fsgroupchangepolicy=always",
			fsGroupChangePolicy: &always,
			setupFunc: func(path string) error {
				info, err := os.Lstat(path)
				if err != nil {
					return err
				}
				// change mode of root folder to be right
				err = os.Chmod(path, info.Mode()|expectedMask)
				if err != nil {
					return err
				}

				// create a subdirectory with invalid permissions
				rogueDir := filepath.Join(path, "roguedir")
				nosgidPerm := info.Mode() &^ os.ModeSetgid
				err = os.Mkdir(rogueDir, nosgidPerm)
				if err != nil {
					return err
				}
				return nil
			},
			assertFunc: func(path string) error {
				rogueDir := filepath.Join(path, "roguedir")
				hasCorrectPermissions := verifyDirectoryPermission(rogueDir, false /*readOnly*/)
				if !hasCorrectPermissions {
					return fmt.Errorf("invalid permissions on %s", rogueDir)
				}
				return nil
			},
		},
		{
			description:         "featuregate=on, fsgroupchangepolicy=onrootmismatch,rootdir=validperm",
			fsGroupChangePolicy: &onrootMismatch,
			setupFunc: func(path string) error {
				info, err := os.Lstat(path)
				if err != nil {
					return err
				}
				// change mode of root folder to be right
				err = os.Chmod(path, info.Mode()|expectedMask)
				if err != nil {
					return err
				}

				// create a subdirectory with invalid permissions
				rogueDir := filepath.Join(path, "roguedir")
				err = os.Mkdir(rogueDir, rwMask)
				if err != nil {
					return err
				}
				return nil
			},
			assertFunc: func(path string) error {
				rogueDir := filepath.Join(path, "roguedir")
				hasCorrectPermissions := verifyDirectoryPermission(rogueDir, false /*readOnly*/)
				if hasCorrectPermissions {
					return fmt.Errorf("invalid permissions on %s", rogueDir)
				}
				return nil
			},
		},
		{
			description:         "featuregate=on, fsgroupchangepolicy=onrootmismatch,rootdir=invalidperm",
			fsGroupChangePolicy: &onrootMismatch,
			setupFunc: func(path string) error {
				// change mode of root folder to be right
				err := os.Chmod(path, 0770)
				if err != nil {
					return err
				}

				// create a subdirectory with invalid permissions
				rogueDir := filepath.Join(path, "roguedir")
				err = os.Mkdir(rogueDir, rwMask)
				if err != nil {
					return err
				}
				return nil
			},
			assertFunc: func(path string) error {
				rogueDir := filepath.Join(path, "roguedir")
				hasCorrectPermissions := verifyDirectoryPermission(rogueDir, false /*readOnly*/)
				if !hasCorrectPermissions {
					return fmt.Errorf("invalid permissions on %s", rogueDir)
				}
				return nil
			},
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			tmpDir, err := utiltesting.MkTmpdir("volume_linux_ownership")
			if err != nil {
				t.Fatalf("error creating temp dir: %v", err)
			}

			defer os.RemoveAll(tmpDir)
			info, err := os.Lstat(tmpDir)
			if err != nil {
				t.Fatalf("error reading permission of tmpdir: %v", err)
			}

			stat, ok := info.Sys().(*syscall.Stat_t)
			if !ok || stat == nil {
				t.Fatalf("error reading permission stats for tmpdir: %s", tmpDir)
			}

			var expectedGid int64 = int64(stat.Gid)
			err = test.setupFunc(tmpDir)
			if err != nil {
				t.Errorf("for %s error running setup with: %v", test.description, err)
			}

			mounter := &localFakeMounter{path: tmpDir}
			err = SetVolumeOwnership(mounter, &expectedGid, test.fsGroupChangePolicy, nil)
			if err != nil {
				t.Errorf("for %s error changing ownership with: %v", test.description, err)
			}
			err = test.assertFunc(tmpDir)
			if err != nil {
				t.Errorf("for %s error verifying permissions with: %v", test.description, err)
			}
		})
	}
}

// verifyDirectoryPermission checks if given path has directory permissions
// that is expected by k8s. If returns true if it does otherwise false
func verifyDirectoryPermission(path string, readonly bool) bool {
	info, err := os.Lstat(path)
	if err != nil {
		return false
	}
	stat, ok := info.Sys().(*syscall.Stat_t)
	if !ok || stat == nil {
		return false
	}
	unixPerms := rwMask

	if readonly {
		unixPerms = roMask
	}

	unixPerms |= execMask
	filePerm := info.Mode().Perm()
	if (unixPerms&filePerm == unixPerms) && (info.Mode()&os.ModeSetgid != 0) {
		return true
	}
	return false
}

func TestSetVolumeOwnershipOwner(t *testing.T) {
	fsGroup := int64(3000)
	currentUid := os.Geteuid()
	if currentUid != 0 {
		t.Skip("running as non-root")
	}
	currentGid := os.Getgid()

	tests := []struct {
		description string
		fsGroup     *int64
		setupFunc   func(path string) error
		assertFunc  func(path string) error
	}{
		{
			description: "fsGroup=nil",
			fsGroup:     nil,
			setupFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0755)
				if err != nil {
					return err
				}
				file.Close()
				return nil
			},
			assertFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				if !verifyFileOwner(filename, currentUid, currentGid) {
					return fmt.Errorf("invalid owner on %s", filename)
				}
				return nil
			},
		},
		{
			description: "*fsGroup=3000",
			fsGroup:     &fsGroup,
			setupFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0755)
				if err != nil {
					return err
				}
				file.Close()
				return nil
			},
			assertFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				if !verifyFileOwner(filename, currentUid, int(fsGroup)) {
					return fmt.Errorf("invalid owner on %s", filename)
				}
				return nil
			},
		},
		{
			description: "symlink",
			fsGroup:     &fsGroup,
			setupFunc: func(path string) error {
				filename := filepath.Join(path, "file.txt")
				file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0755)
				if err != nil {
					return err
				}
				file.Close()

				symname := filepath.Join(path, "file_link.txt")
				err = os.Symlink(filename, symname)
				if err != nil {
					return err
				}

				return nil
			},
			assertFunc: func(path string) error {
				symname := filepath.Join(path, "file_link.txt")
				if !verifyFileOwner(symname, currentUid, int(fsGroup)) {
					return fmt.Errorf("invalid owner on %s", symname)
				}
				return nil
			},
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			tmpDir, err := utiltesting.MkTmpdir("volume_linux_ownership")
			if err != nil {
				t.Fatalf("error creating temp dir: %v", err)
			}

			defer os.RemoveAll(tmpDir)

			err = test.setupFunc(tmpDir)
			if err != nil {
				t.Errorf("for %s error running setup with: %v", test.description, err)
			}

			mounter := &localFakeMounter{path: tmpDir}
			always := v1.FSGroupChangeAlways
			err = SetVolumeOwnership(mounter, test.fsGroup, &always, nil)
			if err != nil {
				t.Errorf("for %s error changing ownership with: %v", test.description, err)
			}
			err = test.assertFunc(tmpDir)
			if err != nil {
				t.Errorf("for %s error verifying permissions with: %v", test.description, err)
			}
		})
	}
}

// verifyFileOwner checks if given path is owned by uid and gid.
// It returns true if it is otherwise false.
func verifyFileOwner(path string, uid, gid int) bool {
	info, err := os.Lstat(path)
	if err != nil {
		return false
	}
	stat, ok := info.Sys().(*syscall.Stat_t)
	if !ok || stat == nil {
		return false
	}

	if int(stat.Uid) != uid || int(stat.Gid) != gid {
		return false
	}

	return true
}
<|endoftext|>


<commit_before>package main

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

	"github.com/pkg/errors"
	"github.com/yamnikov-oleg/avamon-bot/monitor"
)

type quotedSplitter struct {
	sections []string
	buffer   []rune
	quote    *rune
}

func (s *quotedSplitter) next(rn rune) {
	switch {
	case s.quote != nil && rn == *s.quote:
		s.quote = nil
	case s.quote == nil && (rn == '\'' || rn == '"'):
		s.quote = &rn
	case s.quote == nil && rn == ' ':
		if s.buffer != nil {
			s.sections = append(s.sections, string(s.buffer))
			s.buffer = nil
		}
	default:
		s.buffer = append(s.buffer, rn)
	}
}

func (s *quotedSplitter) finish() []string {
	if s.buffer != nil {
		s.sections = append(s.sections, string(s.buffer))
	}
	return s.sections
}

func splitWithQuotes(s string) []string {
	s = strings.TrimSpace(s)
	qs := quotedSplitter{}
	rns := []rune(s)
	for _, rn := range rns {
		qs.next(rn)
	}
	return qs.finish()
}

func parseTarget(args []string) (monitor.Target, error) {
	if len(args) != 3 {
		return monitor.Target{}, errors.New("Target parsing requires 3 arguments")
	}

	var target monitor.Target

	_, err := fmt.Sscan(args[0], &target.ID)
	if err != nil {
		return monitor.Target{}, errors.Wrap(err, "could not parse target id")
	}

	target.Title = args[1]
	target.URL = args[2]

	return target, nil
}

func parseStatus(args []string) (monitor.Status, error) {
	if len(args) != 4 {
		return monitor.Status{}, errors.New("Status parsing requires 4 arguments")
	}

	var status monitor.Status

	stype, ok := monitor.ScanStatusTypeSoft(args[0])
	if !ok {
		return monitor.Status{}, fmt.Errorf("Unknown status type: %v", args[0])
	}
	status.Type = stype

	status.Err = errors.New(args[1])

	rt, err := time.ParseDuration(args[2])
	if err != nil {
		return monitor.Status{}, errors.Wrap(err, "could not parse response time")
	}
	status.ResponseTime = rt

	_, err = fmt.Sscan(args[3], &status.HTTPStatusCode)
	if err != nil {
		return monitor.Status{}, errors.Wrap(err, "could not parse http status code")
	}

	return status, nil
}

func main() {
	host := flag.String("h", "localhost", "Host of redis server")
	port := flag.Uint("p", 6379, "Port of redis server")
	pwd := flag.String("pwd", "", "Password of redis server")
	db := flag.Int("db", 0, "Database to select")

	flag.Parse()

	// targets := monitor.TargetsSlice{}
	rs := monitor.NewRedisStore(monitor.RedisOptions{
		Host:     *host,
		Port:     *port,
		Password: *pwd,
		DB:       *db,
	})

	if err := rs.Ping(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	fmt.Println("Connected to redis.")

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("> ")
		line, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		cmdline := splitWithQuotes(line)
		if len(cmdline) == 0 {
			continue
		}

		cmd := cmdline[0]
		args := cmdline[1:]
		switch cmd {
		case "help":
			fmt.Println("Usage:")
			fmt.Println("  help      - Print help.")
			fmt.Println("  ping      - Ping redis server.")
			fmt.Println("  exit      - Exit this CLI.")
			fmt.Println("  scan      - Scan redis database for statuses.")
			fmt.Println("  getstatus - Get status of a target.")
			fmt.Println("              Arguments:")
			fmt.Println("                <tid>   - Target ID")
			fmt.Println("                <title> - Target Title")
			fmt.Println("                <url>   - Target URL")
			fmt.Println("  setstatus - Set status of a target.")
			fmt.Println("              Arguments:")
			fmt.Println("                <tid>   - Target ID")
			fmt.Println("                <title> - Target Title")
			fmt.Println("                <url>   - Target URL")
			fmt.Println("                <type>  - Status type (e.g. 'ok' or 'generic error')")
			fmt.Println("                <err>   - Error message")
			fmt.Println("                <time>  - Response time (e.g. 200ms)")
			fmt.Println("                <http>  - HTTP status code (e.g. 200) or 0")
			fmt.Println("                <exp>   - Expiration time (e.g. 5s)")

		case "exit":
			fmt.Println("Bye")
			os.Exit(0)

		case "ping":
			err := rs.Ping()
			if err != nil {
				fmt.Println(err)
				continue
			}
			fmt.Println("Pong")

		case "scan":
			tarstats, err := rs.Scan()
			if err != nil {
				fmt.Println(err)
				continue
			}

			fmt.Printf("Found %v statuses\n", len(tarstats))
			for _, ts := range tarstats {
				fmt.Println(ts)
			}

		case "getstatus":
			if len(args) != 3 {
				fmt.Println("This command requires exactly 3 arguments:")
				fmt.Println("getstatus <tid> <title> <url>")
				continue
			}

			target, err := parseTarget(args[0:3])
			if err != nil {
				fmt.Println(err)
				continue
			}

			status, ok, err := rs.GetStatus(target)
			if err != nil {
				fmt.Println(err)
				continue
			}
			if !ok {
				fmt.Println("N/A")
				continue
			}
			fmt.Println(status.ExpandedString())

		case "setstatus":
			if len(args) != 8 {
				fmt.Println("This command requires exactly 8 arguments:")
				fmt.Println("setstatus <tid> <title> <url> <type> <err> <time> <http> <exp>")
				continue
			}

			target, err := parseTarget(args[0:3])
			if err != nil {
				fmt.Println(err)
				continue
			}
			status, err := parseStatus(args[3:7])
			if err != nil {
				fmt.Println(err)
				continue
			}

			exp, err := time.ParseDuration(args[7])
			if err != nil {
				fmt.Printf("Could not parse expiration time: %v\n", err)
				continue
			}

			err = rs.SetStatus(target, status, exp)
			if err != nil {
				fmt.Println(err)
				continue
			}

		default:
			fmt.Println("Unknown command. Print 'help' for list of available commands.")
		}
	}
}
<commit_msg>:lipstick: Make avamon-redisstore CLI rerun last command on empty input<commit_after>package main

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

	"github.com/pkg/errors"
	"github.com/yamnikov-oleg/avamon-bot/monitor"
)

type quotedSplitter struct {
	sections []string
	buffer   []rune
	quote    *rune
}

func (s *quotedSplitter) next(rn rune) {
	switch {
	case s.quote != nil && rn == *s.quote:
		s.quote = nil
	case s.quote == nil && (rn == '\'' || rn == '"'):
		s.quote = &rn
	case s.quote == nil && rn == ' ':
		if s.buffer != nil {
			s.sections = append(s.sections, string(s.buffer))
			s.buffer = nil
		}
	default:
		s.buffer = append(s.buffer, rn)
	}
}

func (s *quotedSplitter) finish() []string {
	if s.buffer != nil {
		s.sections = append(s.sections, string(s.buffer))
	}
	return s.sections
}

func splitWithQuotes(s string) []string {
	s = strings.TrimSpace(s)
	qs := quotedSplitter{}
	rns := []rune(s)
	for _, rn := range rns {
		qs.next(rn)
	}
	return qs.finish()
}

func parseTarget(args []string) (monitor.Target, error) {
	if len(args) != 3 {
		return monitor.Target{}, errors.New("Target parsing requires 3 arguments")
	}

	var target monitor.Target

	_, err := fmt.Sscan(args[0], &target.ID)
	if err != nil {
		return monitor.Target{}, errors.Wrap(err, "could not parse target id")
	}

	target.Title = args[1]
	target.URL = args[2]

	return target, nil
}

func parseStatus(args []string) (monitor.Status, error) {
	if len(args) != 4 {
		return monitor.Status{}, errors.New("Status parsing requires 4 arguments")
	}

	var status monitor.Status

	stype, ok := monitor.ScanStatusTypeSoft(args[0])
	if !ok {
		return monitor.Status{}, fmt.Errorf("Unknown status type: %v", args[0])
	}
	status.Type = stype

	status.Err = errors.New(args[1])

	rt, err := time.ParseDuration(args[2])
	if err != nil {
		return monitor.Status{}, errors.Wrap(err, "could not parse response time")
	}
	status.ResponseTime = rt

	_, err = fmt.Sscan(args[3], &status.HTTPStatusCode)
	if err != nil {
		return monitor.Status{}, errors.Wrap(err, "could not parse http status code")
	}

	return status, nil
}

func main() {
	host := flag.String("h", "localhost", "Host of redis server")
	port := flag.Uint("p", 6379, "Port of redis server")
	pwd := flag.String("pwd", "", "Password of redis server")
	db := flag.Int("db", 0, "Database to select")

	flag.Parse()

	// targets := monitor.TargetsSlice{}
	rs := monitor.NewRedisStore(monitor.RedisOptions{
		Host:     *host,
		Port:     *port,
		Password: *pwd,
		DB:       *db,
	})

	if err := rs.Ping(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	fmt.Println("Connected to redis.")

	var lastCmdline []string
	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("> ")
		line, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		cmdline := splitWithQuotes(line)
		if len(cmdline) == 0 {
			cmdline = lastCmdline
		}
		lastCmdline = cmdline

		cmd := cmdline[0]
		args := cmdline[1:]
		switch cmd {
		case "help":
			fmt.Println("Usage:")
			fmt.Println("  help      - Print help.")
			fmt.Println("  ping      - Ping redis server.")
			fmt.Println("  exit      - Exit this CLI.")
			fmt.Println("  scan      - Scan redis database for statuses.")
			fmt.Println("  getstatus - Get status of a target.")
			fmt.Println("              Arguments:")
			fmt.Println("                <tid>   - Target ID")
			fmt.Println("                <title> - Target Title")
			fmt.Println("                <url>   - Target URL")
			fmt.Println("  setstatus - Set status of a target.")
			fmt.Println("              Arguments:")
			fmt.Println("                <tid>   - Target ID")
			fmt.Println("                <title> - Target Title")
			fmt.Println("                <url>   - Target URL")
			fmt.Println("                <type>  - Status type (e.g. 'ok' or 'generic error')")
			fmt.Println("                <err>   - Error message")
			fmt.Println("                <time>  - Response time (e.g. 200ms)")
			fmt.Println("                <http>  - HTTP status code (e.g. 200) or 0")
			fmt.Println("                <exp>   - Expiration time (e.g. 5s)")

		case "exit":
			fmt.Println("Bye")
			os.Exit(0)

		case "ping":
			err := rs.Ping()
			if err != nil {
				fmt.Println(err)
				continue
			}
			fmt.Println("Pong")

		case "scan":
			tarstats, err := rs.Scan()
			if err != nil {
				fmt.Println(err)
				continue
			}

			fmt.Printf("Found %v statuses\n", len(tarstats))
			for _, ts := range tarstats {
				fmt.Println(ts)
			}

		case "getstatus":
			if len(args) != 3 {
				fmt.Println("This command requires exactly 3 arguments:")
				fmt.Println("getstatus <tid> <title> <url>")
				continue
			}

			target, err := parseTarget(args[0:3])
			if err != nil {
				fmt.Println(err)
				continue
			}

			status, ok, err := rs.GetStatus(target)
			if err != nil {
				fmt.Println(err)
				continue
			}
			if !ok {
				fmt.Println("N/A")
				continue
			}
			fmt.Println(status.ExpandedString())

		case "setstatus":
			if len(args) != 8 {
				fmt.Println("This command requires exactly 8 arguments:")
				fmt.Println("setstatus <tid> <title> <url> <type> <err> <time> <http> <exp>")
				continue
			}

			target, err := parseTarget(args[0:3])
			if err != nil {
				fmt.Println(err)
				continue
			}
			status, err := parseStatus(args[3:7])
			if err != nil {
				fmt.Println(err)
				continue
			}

			exp, err := time.ParseDuration(args[7])
			if err != nil {
				fmt.Printf("Could not parse expiration time: %v\n", err)
				continue
			}

			err = rs.SetStatus(target, status, exp)
			if err != nil {
				fmt.Println(err)
				continue
			}

		default:
			fmt.Println("Unknown command. Print 'help' for list of available commands.")
		}
	}
}
<|endoftext|>


<commit_before>package clients

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/fweikert/continuous-integration/metrics/data"

	"github.com/fweikert/go-buildkite/buildkite"
)

type BuildkiteClient interface {
	GetMostRecentBuilds(*data.PipelineID, int) ([]buildkite.Build, error)
	GetAgents(string) ([]buildkite.Agent, error)
}

type cacheEntry struct {
	sampled time.Time
	values  []interface{}
}

type SimpleBuildkiteClient struct {
	client *buildkite.Client

	mu           sync.Mutex
	cacheTimeout time.Duration
	agentCache   map[string]cacheEntry
	buildCache   map[string]cacheEntry
}

func CreateBuildkiteClient(apiToken string, debug bool, cacheTimeout time.Duration) (BuildkiteClient, error) {
	tokenConfig, err := buildkite.NewTokenConfig(apiToken, debug)
	if err != nil {
		return nil, fmt.Errorf("Could not create Buildkite config: %v", err)
	}
	return &SimpleBuildkiteClient{
			client:       buildkite.NewClient(tokenConfig.Client()),
			cacheTimeout: cacheTimeout,
			agentCache:   make(map[string]cacheEntry),
			buildCache:   make(map[string]cacheEntry),
		},
		nil
}

func (client *SimpleBuildkiteClient) GetMostRecentBuilds(pipeline *data.PipelineID, atLeastNBuilds int) ([]buildkite.Build, error) {
	var listFunc func(opt *buildkite.BuildsListOptions) ([]buildkite.Build, *buildkite.Response, error)
	if pipeline.Slug == "all" {
		listFunc = func(opt *buildkite.BuildsListOptions) ([]buildkite.Build, *buildkite.Response, error) {
			return client.client.Builds.ListByOrg(pipeline.Org, opt)
		}
	} else {
		listFunc = func(opt *buildkite.BuildsListOptions) ([]buildkite.Build, *buildkite.Response, error) {
			return client.client.Builds.ListByPipeline(pipeline.Org, pipeline.Slug, opt)
		}
	}

	wrapperFunc := func(opt buildkite.ListOptions) ([]interface{}, *buildkite.Response, error) {
		buildOpt := buildkite.BuildsListOptions{ListOptions: opt}
		builds, response, err := listFunc(&buildOpt)
		interfaces := make([]interface{}, len(builds))
		for i, b := range builds {
			interfaces[i] = b
		}
		return interfaces, response, err
	}

	results, err := client.getResults(wrapperFunc, atLeastNBuilds, client.buildCache, pipeline.String())
	if err != nil {
		return nil, fmt.Errorf("Failed to retrieve builds for pipeline %s: %v", pipeline, err)
	}

	builds := make([]buildkite.Build, len(results))
	for i, r := range results {
		builds[i] = r.(buildkite.Build)
	}
	return builds, nil
}

func (client *SimpleBuildkiteClient) GetAgents(org string) ([]buildkite.Agent, error) {
	list := func(opt buildkite.ListOptions) ([]interface{}, *buildkite.Response, error) {
		agentOpt := buildkite.AgentListOptions{ListOptions: opt}
		agents, response, err := client.client.Agents.List(org, &agentOpt)
		interfaces := make([]interface{}, len(agents))
		for i, a := range agents {
			interfaces[i] = a
		}
		return interfaces, response, err
	}

	results, err := client.getResults(list, -1, client.agentCache, org)
	if err != nil {
		return nil, fmt.Errorf("Failed to retrieve agents: %v", err)
	}

	agents := make([]buildkite.Agent, len(results))
	for i, r := range results {
		agents[i] = r.(buildkite.Agent)
	}
	return agents, nil
}

func (client *SimpleBuildkiteClient) getResults(listFunc func(opt buildkite.ListOptions) ([]interface{}, *buildkite.Response, error), lastN int, cache map[string]cacheEntry, cacheKey string) ([]interface{}, error) {
	client.mu.Lock()
	defer client.mu.Unlock()
	if entry, ok := cache[cacheKey]; ok {
		if time.Since(entry.sampled) <= client.cacheTimeout && len(entry.values) >= lastN {
			return entry.values[:lastN], nil
		}
		delete(cache, cacheKey)
	}

	results, err := client.getUncachedResults(listFunc, lastN)
	if err != nil {
		return nil, err
	}
	cache[cacheKey] = cacheEntry{
		values:  results,
		sampled: time.Now(),
	}
	return results, nil
}

func (client *SimpleBuildkiteClient) getUncachedResults(listFunc func(opt buildkite.ListOptions) ([]interface{}, *buildkite.Response, error), lastN int) ([]interface{}, error) {
	all_results := make([]interface{}, 0)
	perPage := 100
	if 0 < lastN && lastN < perPage {
		perPage = lastN
	}
	opt := buildkite.ListOptions{Page: 1, PerPage: perPage}
	currPage := 1
	lastPage := 1

	for currPage <= lastPage {
		log.Printf("Retrieving page %d from Buildkite (last=%d).\n", currPage, lastPage)
		results, response, err := listFunc(opt)
		if err != nil {
			return nil, fmt.Errorf("Could not get page %d: %v", currPage, err)
		}

		all_results = append(all_results, results...)
		currPage += 1
		opt.Page = currPage
		lastPage = response.LastPage

		if lastN > -1 && len(all_results) >= lastN {
			break
		}
	}

	if 0 < lastN && lastN < len(all_results) {
		all_results = all_results[:lastN]
	}

	return all_results, nil
}
<commit_msg>Improve logging output for Buildkite requests<commit_after>package clients

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/fweikert/continuous-integration/metrics/data"

	"github.com/fweikert/go-buildkite/buildkite"
)

type BuildkiteClient interface {
	GetMostRecentBuilds(*data.PipelineID, int) ([]buildkite.Build, error)
	GetAgents(string) ([]buildkite.Agent, error)
}

type cacheEntry struct {
	sampled time.Time
	values  []interface{}
}

type SimpleBuildkiteClient struct {
	client *buildkite.Client

	mu           sync.Mutex
	cacheTimeout time.Duration
	agentCache   map[string]cacheEntry
	buildCache   map[string]cacheEntry
}

func CreateBuildkiteClient(apiToken string, debug bool, cacheTimeout time.Duration) (BuildkiteClient, error) {
	tokenConfig, err := buildkite.NewTokenConfig(apiToken, debug)
	if err != nil {
		return nil, fmt.Errorf("Could not create Buildkite config: %v", err)
	}
	return &SimpleBuildkiteClient{
			client:       buildkite.NewClient(tokenConfig.Client()),
			cacheTimeout: cacheTimeout,
			agentCache:   make(map[string]cacheEntry),
			buildCache:   make(map[string]cacheEntry),
		},
		nil
}

func (client *SimpleBuildkiteClient) GetMostRecentBuilds(pipeline *data.PipelineID, atLeastNBuilds int) ([]buildkite.Build, error) {
	var listFunc func(opt *buildkite.BuildsListOptions) ([]buildkite.Build, *buildkite.Response, error)
	if pipeline.Slug == "all" {
		listFunc = func(opt *buildkite.BuildsListOptions) ([]buildkite.Build, *buildkite.Response, error) {
			return client.client.Builds.ListByOrg(pipeline.Org, opt)
		}
	} else {
		listFunc = func(opt *buildkite.BuildsListOptions) ([]buildkite.Build, *buildkite.Response, error) {
			return client.client.Builds.ListByPipeline(pipeline.Org, pipeline.Slug, opt)
		}
	}

	wrapperFunc := func(opt buildkite.ListOptions) ([]interface{}, *buildkite.Response, error) {
		buildOpt := buildkite.BuildsListOptions{ListOptions: opt}
		builds, response, err := listFunc(&buildOpt)
		interfaces := make([]interface{}, len(builds))
		for i, b := range builds {
			interfaces[i] = b
		}
		return interfaces, response, err
	}

	results, err := client.getResults(wrapperFunc, atLeastNBuilds, client.buildCache, pipeline.String())
	if err != nil {
		return nil, fmt.Errorf("Failed to retrieve builds for pipeline %s: %v", pipeline, err)
	}

	builds := make([]buildkite.Build, len(results))
	for i, r := range results {
		builds[i] = r.(buildkite.Build)
	}
	return builds, nil
}

func (client *SimpleBuildkiteClient) GetAgents(org string) ([]buildkite.Agent, error) {
	list := func(opt buildkite.ListOptions) ([]interface{}, *buildkite.Response, error) {
		agentOpt := buildkite.AgentListOptions{ListOptions: opt}
		agents, response, err := client.client.Agents.List(org, &agentOpt)
		interfaces := make([]interface{}, len(agents))
		for i, a := range agents {
			interfaces[i] = a
		}
		return interfaces, response, err
	}

	results, err := client.getResults(list, -1, client.agentCache, org)
	if err != nil {
		return nil, fmt.Errorf("Failed to retrieve agents: %v", err)
	}

	agents := make([]buildkite.Agent, len(results))
	for i, r := range results {
		agents[i] = r.(buildkite.Agent)
	}
	return agents, nil
}

func (client *SimpleBuildkiteClient) getResults(listFunc func(opt buildkite.ListOptions) ([]interface{}, *buildkite.Response, error), lastN int, cache map[string]cacheEntry, cacheKey string) ([]interface{}, error) {
	client.mu.Lock()
	defer client.mu.Unlock()
	if entry, ok := cache[cacheKey]; ok {
		if time.Since(entry.sampled) <= client.cacheTimeout && len(entry.values) >= lastN {
			return entry.values[:lastN], nil
		}
		delete(cache, cacheKey)
	}

	results, err := client.getUncachedResults(listFunc, lastN, cacheKey)
	if err != nil {
		return nil, err
	}
	cache[cacheKey] = cacheEntry{
		values:  results,
		sampled: time.Now(),
	}
	return results, nil
}

func (client *SimpleBuildkiteClient) getUncachedResults(listFunc func(opt buildkite.ListOptions) ([]interface{}, *buildkite.Response, error), lastN int, cacheKey string) ([]interface{}, error) {
	all_results := make([]interface{}, 0)
	perPage := 100
	if 0 < lastN && lastN < perPage {
		perPage = lastN
	}
	opt := buildkite.ListOptions{Page: 1, PerPage: perPage}
	currPage := 1
	lastPage := 1

	for currPage <= lastPage {
		log.Printf("Buildkite: Fetching page %d for '%s' (last=%d).\n", currPage, cacheKey, lastPage)
		results, response, err := listFunc(opt)
		if err != nil {
			return nil, fmt.Errorf("Could not get page %d: %v", currPage, err)
		}

		all_results = append(all_results, results...)
		currPage += 1
		opt.Page = currPage
		lastPage = response.LastPage

		if lastN > -1 && len(all_results) >= lastN {
			break
		}
	}

	if 0 < lastN && lastN < len(all_results) {
		all_results = all_results[:lastN]
	}

	return all_results, nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2014-2016 Miquel Sabaté Solà <mikisabate@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package lib

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
)

// Done this way to test it.
var editCommand = func() error {
	cmd := exec.Command(editor())
	cmd.Dir = filepath.Join(home(), dirName, newDir)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	return cmd.Run()
}

// Edit performs the default command. That is, it fetches all the topics, opens
// up the default editor and pushes the changes.
func Edit() error {
	// Fetch the topics from the server.
	if err := fetch(); err != nil {
		return fromError(err)
	}

	// Open up the editor.
	if err := editCommand(); err != nil {
		return fromError(err)
	}

	// Push all the changed files.
	changed := changedTopics()
	if len(changed) > 0 {
		fmt.Printf("Pushing your changes to the server.\n")
		pushTopics(changed)
	}
	return nil
}

// List simply shows the currently available topics.
func List() error {
	// Try to fetch them if no one else has done it. We can safely ignore the
	// error since we can still cache it if it exists. Otherwise it's not such
	// a pain to get an empty list on weird scenarios.
	_ = fetch()

	var topics []Topic
	readTopics(&topics)
	for _, v := range topics {
		fmt.Printf("%v\n", v.Name)
	}
	return nil
}

// Create creates a new topic on the server.
func Create(name string) error {
	// Perform the HTTP request.
	t := &Topic{Name: name}
	body, _ := json.Marshal(t)
	res, err := getResponse("POST", "/topics", bytes.NewReader(body))
	if err != nil {
		return err
	}

	// Parse the newly created topic and add it to the list.
	if err = topicResponse(t, res); err != nil {
		return NewError("could not create this topic")
	}
	addTopic(t)
	return nil
}

// Delete deletes the specified topic from the server.
func Delete(name string) error {
	var topics, actual []Topic
	var id string

	// Get the list of topics straight.
	readTopics(&topics)
	for _, v := range topics {
		if v.Name == name {
			id = v.ID
		} else {
			actual = append(actual, v)
		}
	}
	if id == "" {
		return unknownTopic(name)
	}

	// Perform the HTTP request.
	if _, err := getResponse("DELETE", "/topics/"+id, nil); err != nil {
		return err
	}

	// On the system.
	writeTopics(actual)
	file := filepath.Join(home(), dirName, oldDir, name+".md")
	_ = os.RemoveAll(file)
	file = filepath.Join(home(), dirName, newDir, name+".md")
	_ = os.RemoveAll(file)
	return nil
}

// Rename changes the name of the given topic with the new one.
func Rename(oldName, newName string) error {
	var topics []Topic
	var id string

	readTopics(&topics)
	for k, v := range topics {
		if v.Name == oldName {
			id = v.ID
			topics[k].Name = newName
		}
	}
	if id == "" {
		return unknownTopic(oldName)
	}

	// Perform the HTTP Request.
	t := &Topic{Name: newName}
	body, _ := json.Marshal(t)
	res, err := getResponse("PUT", "/topics/"+id, bytes.NewReader(body))
	if err != nil {
		return err
	}
	if err = topicResponse(t, res); err != nil {
		return NewError("could not rename this topic")
	}

	// Update the system.
	writeTopics(topics)
	file := filepath.Join(home(), dirName, oldDir)
	_ = os.Rename(filepath.Join(file, oldName), filepath.Join(file, newName))
	file = filepath.Join(home(), dirName, newDir)
	_ = os.Rename(filepath.Join(file, oldName), filepath.Join(file, newName))
	return nil
}
<commit_msg>Be more specific on errors on create and delete<commit_after>// Copyright (C) 2014-2016 Miquel Sabaté Solà <mikisabate@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package lib

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
)

// Done this way to test it.
var editCommand = func() error {
	cmd := exec.Command(editor())
	cmd.Dir = filepath.Join(home(), dirName, newDir)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	return cmd.Run()
}

// Edit performs the default command. That is, it fetches all the topics, opens
// up the default editor and pushes the changes.
func Edit() error {
	// Fetch the topics from the server.
	if err := fetch(); err != nil {
		return fromError(err)
	}

	// Open up the editor.
	if err := editCommand(); err != nil {
		return fromError(err)
	}

	// Push all the changed files.
	changed := changedTopics()
	if len(changed) > 0 {
		fmt.Printf("Pushing your changes to the server.\n")
		pushTopics(changed)
	}
	return nil
}

// List simply shows the currently available topics.
func List() error {
	// Try to fetch them if no one else has done it. We can safely ignore the
	// error since we can still cache it if it exists. Otherwise it's not such
	// a pain to get an empty list on weird scenarios.
	_ = fetch()

	var topics []Topic
	readTopics(&topics)
	for _, v := range topics {
		fmt.Printf("%v\n", v.Name)
	}
	return nil
}

// Create creates a new topic on the server.
func Create(name string) error {
	// Perform the HTTP request.
	t := &Topic{Name: name}
	body, _ := json.Marshal(t)
	res, err := getResponse("POST", "/topics", bytes.NewReader(body))
	if err != nil {
		return err
	}

	// Parse the newly created topic and add it to the list.
	if err = topicResponse(t, res); err != nil {
		return NewError("could not create this topic: " + err.Error())
	}
	addTopic(t)
	return nil
}

// Delete deletes the specified topic from the server.
func Delete(name string) error {
	var topics, actual []Topic
	var id string

	// Get the list of topics straight.
	readTopics(&topics)
	for _, v := range topics {
		if v.Name == name {
			id = v.ID
		} else {
			actual = append(actual, v)
		}
	}
	if id == "" {
		return unknownTopic(name)
	}

	// Perform the HTTP request.
	if _, err := getResponse("DELETE", "/topics/"+id, nil); err != nil {
		return err
	}

	// On the system.
	writeTopics(actual)
	file := filepath.Join(home(), dirName, oldDir, name+".md")
	_ = os.RemoveAll(file)
	file = filepath.Join(home(), dirName, newDir, name+".md")
	_ = os.RemoveAll(file)
	return nil
}

// Rename changes the name of the given topic with the new one.
func Rename(oldName, newName string) error {
	var topics []Topic
	var id string

	readTopics(&topics)
	for k, v := range topics {
		if v.Name == oldName {
			id = v.ID
			topics[k].Name = newName
		}
	}
	if id == "" {
		return unknownTopic(oldName)
	}

	// Perform the HTTP Request.
	t := &Topic{Name: newName}
	body, _ := json.Marshal(t)
	res, err := getResponse("PUT", "/topics/"+id, bytes.NewReader(body))
	if err != nil {
		return err
	}
	if err = topicResponse(t, res); err != nil {
		return NewError("could not rename this topic: " + err.Error())
	}

	// Update the system.
	writeTopics(topics)
	file := filepath.Join(home(), dirName, oldDir)
	_ = os.Rename(filepath.Join(file, oldName), filepath.Join(file, newName))
	file = filepath.Join(home(), dirName, newDir)
	_ = os.Rename(filepath.Join(file, oldName), filepath.Join(file, newName))
	return nil
}
<|endoftext|>


<commit_before>// $G $D/$F.go && $L $F.$A && ./$A.out

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

package main

type Box struct {};
var m map[string] *Box;

func main() {
  s := "foo";
  var x *Box = nil;
  m[s] = x;
}

/*
bug058.go:9: illegal types for operand: INDEX
	(MAP[<string>*STRING]*<Box>{})
	(<string>*STRING)
*/
<commit_msg>- added calc cmd<commit_after>// $G $D/$F.go && $L $F.$A && ./$A.out

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

package main

type Box struct {};
var m *map[string] *Box;

func main() {
  m := new(map[string] *Box);
  s := "foo";
  var x *Box = nil;
  m[s] = x;
}

/*
bug058.go:9: illegal types for operand: INDEX
	(MAP[<string>*STRING]*<Box>{})
	(<string>*STRING)
*/
<|endoftext|>


<commit_before>// Written by Charles Cochrane, 2015
// Use of this source code is governed by a MIT license that can be found
// in the LICENSE file.

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"runtime"
	"sync"

	"github.com/boltdb/bolt"
	"github.com/julienschmidt/httprouter"
)

type SurnameStruct struct {
	Surname string        `json:"Surname"`
	Entries []EntryStruct `json:"Entries"`
}

type EntryStruct struct {
	FirstName   string `json:"FirstName"`
	TelNo       string `json:"TelNo"`
	Line1       string `json:"Line1"`
	Line2       string `json:"Line2"`
	TownCity    string `json:"TownCity"`
	CountyState string `json:"CountyState"`
	Country     string `json:"Country"`
	ZipPostal   string `json:"ZipPostal"`
}

var BoltClient struct {
	DB    *bolt.DB
	Mutex *sync.RWMutex
}

func main() {

	SetProc()
	NewBoltClient()
	defer BoltClient.DB.Close()

	//start the server
	log.Println("Listening on port :3000")
	log.Fatal(http.ListenAndServe(":3000", NewRouter()))
}

func listHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
	log.Println("list")

	BoltClient.Mutex.RLock()

	//list all entries in the bucket
	fmt.Println("all in the phonebook:")
	BoltClient.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("phoneBook"))
		b.ForEach(func(k, v []byte) error {
			fmt.Printf("key=%s, value=%s\n", k, v)
			return nil
		})
		return nil
	})

	BoltClient.Mutex.RUnlock()
}

//func searchHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
//	log.Println("search", params.ByName("surname"))
//}

func getEntryHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
	log.Println("get", params.ByName("surname"))

	BoltClient.Mutex.RLock()

	//Get from bucket
	BoltClient.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("phoneBook"))
		v := b.Get([]byte(params.ByName("surname")))
		if v == nil {
			w.WriteHeader(404)
		} else {
			//fmt.Printf("The answer is: %s\n", v)
			fmt.Fprintf(w, "%s\n", v)
		}
		return nil
	})

	BoltClient.Mutex.RUnlock()

}

func putEntryHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
	log.Println("put", params.ByName("surname"))

	//get the content of the HTTP request
	defer r.Body.Close()
	bodyByte, _ := ioutil.ReadAll(r.Body)
	body := string(bodyByte)
	log.Println("Content: " + body)

	//marshal unmarshal Json
	//check if that surname is present
	//if it is then retrieve and add record
	//else

	BoltClient.Mutex.Lock()

	err := BoltClient.DB.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("phoneBook"))
		err := b.Put([]byte(params.ByName("surname")), []byte(body))
		return err
	})

	BoltClient.Mutex.Unlock()

	if err != nil {
		log.Print(err)
	}
}

func delEntryHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
	log.Println("del", params.ByName("surname"))

	BoltClient.Mutex.Lock()

	//Delete from bucket
	err := BoltClient.DB.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("phoneBook"))
		err := b.Delete([]byte(params.ByName("surname")))
		return err
	})

	BoltClient.Mutex.Unlock()

	if err != nil {
		log.Print(err)
	}

}

//NewRouter returns a httprouter.Router complete with the routes
func NewRouter() *httprouter.Router {

	router := httprouter.New()
	router.GET("/list", listHandler)
	//router.GET("/search/:surname", searchHandler) //need to fix this, to /search?surname=bob
	router.GET("/entry/:surname", getEntryHandler) //is this the search
	router.PUT("/entry/:surname", putEntryHandler)
	router.DELETE("/entry/:surname", delEntryHandler)

	return router
}

//SetProc sets the program to use 2 proccessor cores
func SetProc() {
	runtime.GOMAXPROCS(2)
}

//NewBoltClient produces a BoltClient and Mutex lock and assigns them
//to the global BoltClient variable
func NewBoltClient() { //*BoltClient {
	//Open DB
	BoltDB, err := bolt.Open("phoneBook.db", 0600, nil)
	if err != nil {
		log.Fatal(err)
	}

	//open (or create if not present) bucket
	err = BoltDB.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("phoneBook"))

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

	//make a read/write mutex
	mutex := &sync.RWMutex{}

	//assign the pointers to the mutex and DB to the Global
	//BoltClient var
	BoltClient.DB = BoltDB
	BoltClient.Mutex = mutex
}
<commit_msg>list handler added<commit_after>// Written by Charles Cochrane, 2015
// Use of this source code is governed by a MIT license that can be found
// in the LICENSE file.

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"runtime"
	"sync"

	"github.com/boltdb/bolt"
	"github.com/julienschmidt/httprouter"
)

type SurnameStruct struct {
	Surname string        `json:"Surname"`
	Entries []EntryStruct `json:"Entries"`
}

type EntryStruct struct {
	FirstName   string `json:"FirstName"`
	TelNo       string `json:"TelNo"`
	Line1       string `json:"Line1"`
	Line2       string `json:"Line2"`
	TownCity    string `json:"TownCity"`
	CountyState string `json:"CountyState"`
	Country     string `json:"Country"`
	ZipPostal   string `json:"ZipPostal"`
}

var BoltClient struct {
	DB    *bolt.DB
	Mutex *sync.RWMutex
}

func main() {

	SetProc()
	NewBoltClient()
	defer BoltClient.DB.Close()

	//start the server
	log.Println("Listening on port :3000")
	log.Fatal(http.ListenAndServe(":3000", NewRouter()))
}

func listHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
	log.Println("list")

	fmt.Fprint(w, `{"Phone Book":`)
	BoltClient.Mutex.RLock()

	//list all entries in the bucket
	//fmt.Println("all in the phonebook:")
	BoltClient.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("phoneBook"))
		b.ForEach(func(k, v []byte) error {
			//fmt.Printf("key=%s, value=%s\n", k, v)
			fmt.Fprintf(w, "%s,", v) //"%s\n", v)
			return nil
		})
		return nil
	})

	BoltClient.Mutex.RUnlock()
	fmt.Fprint(w, `}`)
}

//func searchHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
//	log.Println("search", params.ByName("surname"))
//}

func getEntryHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
	log.Println("get", params.ByName("surname"))

	BoltClient.Mutex.RLock()

	//Get from bucket
	BoltClient.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("phoneBook"))
		v := b.Get([]byte(params.ByName("surname")))
		if v == nil {
			w.WriteHeader(404)
		} else {
			fmt.Fprintf(w, "%s\n", v)
		}
		return nil
	})

	BoltClient.Mutex.RUnlock()

}

func putEntryHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
	log.Println("put", params.ByName("surname"))

	//get the content of the HTTP request
	defer r.Body.Close()
	bodyByte, _ := ioutil.ReadAll(r.Body)
	body := string(bodyByte)
	log.Println("Content: " + body)

	//marshal unmarshal Json
	//check if that surname is present
	//if it is then retrieve and add record
	//else

	BoltClient.Mutex.Lock()

	err := BoltClient.DB.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("phoneBook"))
		err := b.Put([]byte(params.ByName("surname")), []byte(body))
		return err
	})

	BoltClient.Mutex.Unlock()

	if err != nil {
		log.Print(err)
	}
}

func delEntryHandler(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
	log.Println("del", params.ByName("surname"))

	BoltClient.Mutex.Lock()

	//Delete from bucket
	err := BoltClient.DB.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("phoneBook"))
		err := b.Delete([]byte(params.ByName("surname")))
		return err
	})

	BoltClient.Mutex.Unlock()

	if err != nil {
		log.Print(err)
	}

}

//NewRouter returns a httprouter.Router complete with the routes
func NewRouter() *httprouter.Router {

	router := httprouter.New()
	router.GET("/list", listHandler)
	//router.GET("/search/:surname", searchHandler) //need to fix this, to /search?surname=bob
	router.GET("/entry/:surname", getEntryHandler) //is this the search
	router.PUT("/entry/:surname", putEntryHandler)
	router.DELETE("/entry/:surname", delEntryHandler)

	return router
}

//SetProc sets the program to use 2 proccessor cores
func SetProc() {
	runtime.GOMAXPROCS(2)
}

//NewBoltClient produces a BoltClient and Mutex lock and assigns them
//to the global BoltClient variable
func NewBoltClient() { //*BoltClient {
	//Open DB
	BoltDB, err := bolt.Open("phoneBook.db", 0600, nil)
	if err != nil {
		log.Fatal(err)
	}

	//open (or create if not present) bucket
	err = BoltDB.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("phoneBook"))

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

	//make a read/write mutex
	mutex := &sync.RWMutex{}

	//assign the pointers to the mutex and DB to the Global
	//BoltClient var
	BoltClient.DB = BoltDB
	BoltClient.Mutex = mutex
}
<|endoftext|>


<commit_before>package tweetmap

import (
	"database/sql"
	"errors"
	_ "github.com/lib/pq"
	"os"
)

// Query types
const (
	TWEETS_PAST_24HRS       int = iota
	TWEETS_BY_WARD          int = iota
	HASH_BY_WARD_PAST_24HRS int = iota
)

type Database struct {
	db         *sql.DB
	statements map[int]*sql.Stmt
}

func NewDatabase() (*Database, error) {
	var db Database
	err := db.connect()
	if err != nil {
		return nil, err
	}
	return &db, nil
}

// Returns nil if currently connected to the database
func (Db Database) connected() error {
	if err := Db.db.Ping(); err != nil {
		return err
	}

	return nil
}

// Set up prepared statements for the queries
func (Db *Database) setupStatements() error {
	stmts := make(map[int]string)
	stmts[TWEETS_PAST_24HRS] = "SELECT * FROM \"V_tweetsPast24Hrs\";"
	stmts[TWEETS_BY_WARD] = "SELECT * FROM \"V_tweetsByWardDailyChange\";"

	for k, v := range stmts {
		// Try and prepare the statements
		if p, e := Db.db.Prepare(v); e != nil {
			return e
		} else {
			Db.statements[k] = p
		}
	}
	return nil
}

// Connects to the database
// Sets up all the prepared statements
// Requires the PGUSER, PGPASSWORD, PGHOST, PGPORT, and PGDATABASE environment variables to be set
func (Db *Database) connect() error {
	var connStr string
	connStr = "user=" + os.Getenv("PGUSER") + " "
	connStr += "password=" + os.Getenv("PGPASSWORD") + " "
	connStr += "host=" + os.Getenv("PGHOST") + " "
	connStr += "port=" + os.Getenv("PGPORT") + " "
	connStr += "dbname=" + os.Getenv("PGDATABASE") + " "
	connStr += "search_path=tweetwatcher"

	var err error
	Db.db, err = sql.Open("postgres", connStr)

	if err != nil {
		return err
	}

	if err := Db.connected(); err != nil {
		return nil
	} else {
		return err
	}

	if err := Db.setupStatements(); err != nil {
		Db.Disconnect()
		return err
	}

	return nil
}

// Wrapper for DB.Close()
func (Db *Database) Disconnect() {
	Db.db.Close()
}

// Query
// Returns the result of a query, in marshalled JSON, which then can be served to the client
func (Db *Database) Query(queryType int, args ...interface{}) ([]byte, error) {
	switch queryType {
	case TWEETS_PAST_24HRS:
		return Db.tweetsPast24Hrs()
		break
	}
	return nil, errors.New("Invalid query type")
}

// Returns the total number of tweets per hour
func (Db Database) tweetsPast24Hrs() ([]byte, error) {
	var result string
	err := Db.db.QueryRow("SELECT * FROM \"V_tweetsPast24Hrs\";").Scan(&result)
	if err != nil {
		return nil, err
	}
	return []byte(result), nil
}
<commit_msg>test commit<commit_after>package tweetmap

import (
	"database/sql"
	"errors"
	_ "github.com/lib/pq"
	"os"
)

// Query types
const (
	TWEETS_PAST_24HRS       int = iota
	TWEETS_BY_WARD          int = iota
	HASH_BY_WARD_PAST_24HRS int = iota
)

type Database struct {
	db         *sql.DB
	statements map[int]*sql.Stmt
}

func NewDatabase() (*Database, error) {
	var db Database
	err := db.connect()
	if err != nil {
		return nil, err
	}
	return &db, nil
}

// Returns nil if currently connected to the database
func (Db Database) connected() error {
	if err := Db.db.Ping(); err != nil {
		return err
	}

	return nil
}

// Set up prepared statements for the queries
func (Db *Database) setupStatements() error {
	stmts := make(map[int]string)
	stmts[TWEETS_PAST_24HRS] = "SELECT * FROM \"V_tweetsPast24Hrs\";"
	stmts[TWEETS_BY_WARD] = "SELECT * FROM \"V_tweetsByWardDailyChange\";"

	for k, v := range stmts {
		// Try and prepare the statements
		if p, e := Db.db.Prepare(v); e != nil {
			return e
		} else {
			Db.statements[k] = p
		}
	}
	return nil
}

// Connects to the database
// Sets up all the prepared statements
// Requires the PGUSER, PGPASSWORD, PGHOST, PGPORT, and PGDATABASE environment variables to be set
func (Db *Database) connect() error {
	var connStr string
	connStr = "user=" + os.Getenv("PGUSER") + " "
	connStr += "password=" + os.Getenv("PGPASSWORD") + " "
	connStr += "host=" + os.Getenv("PGHOST") + " "
	connStr += "port=" + os.Getenv("PGPORT") + " "
	connStr += "dbname=" + os.Getenv("PGDATABASE") + " "
	connStr += "search_path=tweetwatcher"

	var err error
	Db.db, err = sql.Open("postgres", connStr)

	if err != nil {
		return err
	}

	if err := Db.connected(); err != nil {
		return nil
	} else {
		return err
	}

	if err := Db.setupStatements(); err != nil {
		Db.Disconnect()
		return err
	}

	return nil
}

// Wrapper for DB.Close()
func (Db *Database) Disconnect() {
	Db.db.Close()
}

// Query
// Returns the result of a query, in marshalled JSON, which then can be served to the client
func (Db *Database) Query(queryType int, args ...interface{}) ([]byte, error) {
	switch queryType {
	case TWEETS_PAST_24HRS:
		return Db.tweetsPast24Hrs()
		break
	}
	return nil, errors.New("Invalid query type")
}

// Returns the total number of tweets per hour
<<<<<<< HEAD
func (Db Database) tweetsPast24Hrs() ([]byte, error) {
=======
func (Db *Database) tweetsPast24Hrs() ([]byte, error) {
>>>>>>> 3092ef1d2fba89209b32d25d117c4fd9915eecef
	var result string
	err := Db.db.QueryRow("SELECT * FROM \"V_tweetsPast24Hrs\";").Scan(&result)
	if err != nil {
		return nil, err
	}
	return []byte(result), nil
}
<|endoftext|>


<commit_before>// Package gc provides garbage collection for go-ipfs.
package gc

import (
	"context"
	"errors"
	"fmt"
	"strings"

	pin "github.com/ipfs/go-ipfs/pin"
	dag "gx/ipfs/QmXv5mwmQ74r4aiHcNeQ4GAmfB3aWJuqaE4WyDfDfvkgLM/go-merkledag"
	bserv "gx/ipfs/Qma2KhbQarYTkmSJAeaMGRAg8HAXAhEWK8ge4SReG7ZSD3/go-blockservice"

	cid "gx/ipfs/QmPSQnBKM9g7BaUcZCvswUJVscQ1ipjmwxN5PXCjkp9EQ7/go-cid"
	logging "gx/ipfs/QmRREK2CAZ5Re2Bd9zZFG6FeYDppUWt5cMgsoUEp3ktgSr/go-log"
	dstore "gx/ipfs/QmSpg1CvpXQQow5ernt1gNBXaXV6yxyNqi7XoeerWfzB5w/go-datastore"
	"gx/ipfs/QmVkMRSkXrpjqrroEXWuYBvDBnXCdMMY6gsKicBGVGUqKT/go-verifcid"
	offline "gx/ipfs/QmcRC35JF2pJQneAxa5LdQBQRumWggccWErogSrCkS1h8T/go-ipfs-exchange-offline"
	ipld "gx/ipfs/QmdDXJs4axxefSPgK6Y1QhpJWKuDPnGJiqgq4uncb4rFHL/go-ipld-format"
	bstore "gx/ipfs/QmegPGspn3RpTMQ23Fd3GVVMopo1zsEMurudbFMZ5UXBLH/go-ipfs-blockstore"
)

var log = logging.Logger("gc")

// Result represents an incremental output from a garbage collection
// run.  It contains either an error, or the cid of a removed object.
type Result struct {
	KeyRemoved cid.Cid
	Error      error
}

// GC performs a mark and sweep garbage collection of the blocks in the blockstore
// first, it creates a 'marked' set and adds to it the following:
// - all recursively pinned blocks, plus all of their descendants (recursively)
// - bestEffortRoots, plus all of its descendants (recursively)
// - all directly pinned blocks
// - all blocks utilized internally by the pinner
//
// The routine then iterates over every block in the blockstore and
// deletes any block that is not found in the marked set.
func GC(ctx context.Context, bs bstore.GCBlockstore, dstor dstore.Datastore, pn pin.Pinner, bestEffortRoots []cid.Cid) <-chan Result {

	elock := log.EventBegin(ctx, "GC.lockWait")
	unlocker := bs.GCLock()
	elock.Done()
	elock = log.EventBegin(ctx, "GC.locked")
	emark := log.EventBegin(ctx, "GC.mark")

	bsrv := bserv.New(bs, offline.Exchange(bs))
	ds := dag.NewDAGService(bsrv)

	output := make(chan Result, 128)

	go func() {
		defer close(output)
		defer unlocker.Unlock()
		defer elock.Done()

		gcs, err := ColoredSet(ctx, pn, ds, bestEffortRoots, output)
		if err != nil {
			output <- Result{Error: err}
			return
		}
		emark.Append(logging.LoggableMap{
			"blackSetSize": fmt.Sprintf("%d", gcs.Len()),
		})
		emark.Done()
		esweep := log.EventBegin(ctx, "GC.sweep")

		keychan, err := bs.AllKeysChan(ctx)
		if err != nil {
			output <- Result{Error: err}
			return
		}

		errors := false
		var removed uint64

	loop:
		for {
			select {
			case k, ok := <-keychan:
				if !ok {
					break loop
				}
				if !gcs.Has(k) {
					err := bs.DeleteBlock(k)
					removed++
					if err != nil {
						errors = true
						output <- Result{Error: &CannotDeleteBlockError{k, err}}
						//log.Errorf("Error removing key from blockstore: %s", err)
						// continue as error is non-fatal
						continue loop
					}
					select {
					case output <- Result{KeyRemoved: k}:
					case <-ctx.Done():
						break loop
					}
				}
			case <-ctx.Done():
				break loop
			}
		}
		esweep.Append(logging.LoggableMap{
			"whiteSetSize": fmt.Sprintf("%d", removed),
		})
		esweep.Done()
		if errors {
			output <- Result{Error: ErrCannotDeleteSomeBlocks}
		}

		defer log.EventBegin(ctx, "GC.datastore").Done()
		gds, ok := dstor.(dstore.GCDatastore)
		if !ok {
			return
		}

		err = gds.CollectGarbage()
		if err != nil {
			output <- Result{Error: err}
			return
		}
	}()

	return output
}

// Descendants recursively finds all the descendants of the given roots and
// adds them to the given cid.Set, using the provided dag.GetLinks function
// to walk the tree.
func Descendants(ctx context.Context, getLinks dag.GetLinks, set *cid.Set, roots []cid.Cid) error {
	verifyGetLinks := func(ctx context.Context, c cid.Cid) ([]*ipld.Link, error) {
		err := verifcid.ValidateCid(c)
		if err != nil {
			return nil, err
		}

		return getLinks(ctx, c)
	}

	verboseCidError := func(err error) error {
		if strings.Contains(err.Error(), verifcid.ErrBelowMinimumHashLength.Error()) ||
			strings.Contains(err.Error(), verifcid.ErrPossiblyInsecureHashFunction.Error()) {
			err = fmt.Errorf("\"%s\"\nPlease run 'ipfs pin verify'"+
				" to list insecure hashes. If you want to read them,"+
				" please downgrade your go-ipfs to 0.4.13\n", err)
			log.Error(err)
		}
		return err
	}

	for _, c := range roots {
		set.Add(c)

		// EnumerateChildren recursively walks the dag and adds the keys to the given set
		err := dag.EnumerateChildren(ctx, verifyGetLinks, c, set.Visit)

		if err != nil {
			err = verboseCidError(err)
			return err
		}
	}

	return nil
}

// ColoredSet computes the set of nodes in the graph that are pinned by the
// pins in the given pinner.
func ColoredSet(ctx context.Context, pn pin.Pinner, ng ipld.NodeGetter, bestEffortRoots []cid.Cid, output chan<- Result) (*cid.Set, error) {
	// KeySet currently implemented in memory, in the future, may be bloom filter or
	// disk backed to conserve memory.
	errors := false
	gcs := cid.NewSet()
	getLinks := func(ctx context.Context, cid cid.Cid) ([]*ipld.Link, error) {
		links, err := ipld.GetLinks(ctx, ng, cid)
		if err != nil {
			errors = true
			output <- Result{Error: &CannotFetchLinksError{cid, err}}
		}
		return links, nil
	}
	err := Descendants(ctx, getLinks, gcs, pn.RecursiveKeys())
	if err != nil {
		errors = true
		output <- Result{Error: err}
	}

	bestEffortGetLinks := func(ctx context.Context, cid cid.Cid) ([]*ipld.Link, error) {
		links, err := ipld.GetLinks(ctx, ng, cid)
		if err != nil && err != ipld.ErrNotFound {
			errors = true
			output <- Result{Error: &CannotFetchLinksError{cid, err}}
		}
		return links, nil
	}
	err = Descendants(ctx, bestEffortGetLinks, gcs, bestEffortRoots)
	if err != nil {
		errors = true
		output <- Result{Error: err}
	}

	for _, k := range pn.DirectKeys() {
		gcs.Add(k)
	}

	err = Descendants(ctx, getLinks, gcs, pn.InternalPins())
	if err != nil {
		errors = true
		output <- Result{Error: err}
	}

	if errors {
		return nil, ErrCannotFetchAllLinks
	}

	return gcs, nil
}

// ErrCannotFetchAllLinks is returned as the last Result in the GC output
// channel when there was a error creating the marked set because of a
// problem when finding descendants.
var ErrCannotFetchAllLinks = errors.New("garbage collection aborted: could not retrieve some links")

// ErrCannotDeleteSomeBlocks is returned when removing blocks marked for
// deletion fails as the last Result in GC output channel.
var ErrCannotDeleteSomeBlocks = errors.New("garbage collection incomplete: could not delete some blocks")

// CannotFetchLinksError provides detailed information about which links
// could not be fetched and can appear as a Result in the GC output channel.
type CannotFetchLinksError struct {
	Key cid.Cid
	Err error
}

// Error implements the error interface for this type with a useful
// message.
func (e *CannotFetchLinksError) Error() string {
	return fmt.Sprintf("could not retrieve links for %s: %s", e.Key, e.Err)
}

// CannotDeleteBlockError provides detailed information about which
// blocks could not be deleted and can appear as a Result in the GC output
// channel.
type CannotDeleteBlockError struct {
	Key cid.Cid
	Err error
}

// Error implements the error interface for this type with a
// useful message.
func (e *CannotDeleteBlockError) Error() string {
	return fmt.Sprintf("could not remove %s: %s", e.Key, e.Err)
}
<commit_msg>fix(pin): goroutine leaks<commit_after>// Package gc provides garbage collection for go-ipfs.
package gc

import (
	"context"
	"errors"
	"fmt"
	"strings"

	pin "github.com/ipfs/go-ipfs/pin"
	dag "gx/ipfs/QmXv5mwmQ74r4aiHcNeQ4GAmfB3aWJuqaE4WyDfDfvkgLM/go-merkledag"
	bserv "gx/ipfs/Qma2KhbQarYTkmSJAeaMGRAg8HAXAhEWK8ge4SReG7ZSD3/go-blockservice"

	cid "gx/ipfs/QmPSQnBKM9g7BaUcZCvswUJVscQ1ipjmwxN5PXCjkp9EQ7/go-cid"
	logging "gx/ipfs/QmRREK2CAZ5Re2Bd9zZFG6FeYDppUWt5cMgsoUEp3ktgSr/go-log"
	dstore "gx/ipfs/QmSpg1CvpXQQow5ernt1gNBXaXV6yxyNqi7XoeerWfzB5w/go-datastore"
	"gx/ipfs/QmVkMRSkXrpjqrroEXWuYBvDBnXCdMMY6gsKicBGVGUqKT/go-verifcid"
	offline "gx/ipfs/QmcRC35JF2pJQneAxa5LdQBQRumWggccWErogSrCkS1h8T/go-ipfs-exchange-offline"
	ipld "gx/ipfs/QmdDXJs4axxefSPgK6Y1QhpJWKuDPnGJiqgq4uncb4rFHL/go-ipld-format"
	bstore "gx/ipfs/QmegPGspn3RpTMQ23Fd3GVVMopo1zsEMurudbFMZ5UXBLH/go-ipfs-blockstore"
)

var log = logging.Logger("gc")

// Result represents an incremental output from a garbage collection
// run.  It contains either an error, or the cid of a removed object.
type Result struct {
	KeyRemoved cid.Cid
	Error      error
}

// GC performs a mark and sweep garbage collection of the blocks in the blockstore
// first, it creates a 'marked' set and adds to it the following:
// - all recursively pinned blocks, plus all of their descendants (recursively)
// - bestEffortRoots, plus all of its descendants (recursively)
// - all directly pinned blocks
// - all blocks utilized internally by the pinner
//
// The routine then iterates over every block in the blockstore and
// deletes any block that is not found in the marked set.
func GC(ctx context.Context, bs bstore.GCBlockstore, dstor dstore.Datastore, pn pin.Pinner, bestEffortRoots []cid.Cid) <-chan Result {

	elock := log.EventBegin(ctx, "GC.lockWait")
	unlocker := bs.GCLock()
	elock.Done()
	elock = log.EventBegin(ctx, "GC.locked")
	emark := log.EventBegin(ctx, "GC.mark")

	bsrv := bserv.New(bs, offline.Exchange(bs))
	ds := dag.NewDAGService(bsrv)

	output := make(chan Result, 128)

	go func() {
		defer close(output)
		defer unlocker.Unlock()
		defer elock.Done()

		gcs, err := ColoredSet(ctx, pn, ds, bestEffortRoots, output)
		if err != nil {
			select {
			case output <- Result{Error: err}:
			case <-ctx.Done():
			}
			return
		}
		emark.Append(logging.LoggableMap{
			"blackSetSize": fmt.Sprintf("%d", gcs.Len()),
		})
		emark.Done()
		esweep := log.EventBegin(ctx, "GC.sweep")

		keychan, err := bs.AllKeysChan(ctx)
		if err != nil {
			select {
			case output <- Result{Error: err}:
			case <-ctx.Done():
			}
			return
		}

		errors := false
		var removed uint64

	loop:
		for {
			select {
			case k, ok := <-keychan:
				if !ok {
					break loop
				}
				if !gcs.Has(k) {
					err := bs.DeleteBlock(k)
					removed++
					if err != nil {
						errors = true
						output <- Result{Error: &CannotDeleteBlockError{k, err}}
						//log.Errorf("Error removing key from blockstore: %s", err)
						// continue as error is non-fatal
						continue loop
					}
					select {
					case output <- Result{KeyRemoved: k}:
					case <-ctx.Done():
						break loop
					}
				}
			case <-ctx.Done():
				break loop
			}
		}
		esweep.Append(logging.LoggableMap{
			"whiteSetSize": fmt.Sprintf("%d", removed),
		})
		esweep.Done()
		if errors {
			select {
			case output <- Result{Error: ErrCannotDeleteSomeBlocks}:
			case <-ctx.Done():
				return
			}
		}

		defer log.EventBegin(ctx, "GC.datastore").Done()
		gds, ok := dstor.(dstore.GCDatastore)
		if !ok {
			return
		}

		err = gds.CollectGarbage()
		if err != nil {
			select {
			case output <- Result{Error: err}:
			case <-ctx.Done():
			}
			return
		}
	}()

	return output
}

// Descendants recursively finds all the descendants of the given roots and
// adds them to the given cid.Set, using the provided dag.GetLinks function
// to walk the tree.
func Descendants(ctx context.Context, getLinks dag.GetLinks, set *cid.Set, roots []cid.Cid) error {
	verifyGetLinks := func(ctx context.Context, c cid.Cid) ([]*ipld.Link, error) {
		err := verifcid.ValidateCid(c)
		if err != nil {
			return nil, err
		}

		return getLinks(ctx, c)
	}

	verboseCidError := func(err error) error {
		if strings.Contains(err.Error(), verifcid.ErrBelowMinimumHashLength.Error()) ||
			strings.Contains(err.Error(), verifcid.ErrPossiblyInsecureHashFunction.Error()) {
			err = fmt.Errorf("\"%s\"\nPlease run 'ipfs pin verify'"+
				" to list insecure hashes. If you want to read them,"+
				" please downgrade your go-ipfs to 0.4.13\n", err)
			log.Error(err)
		}
		return err
	}

	for _, c := range roots {
		set.Add(c)

		// EnumerateChildren recursively walks the dag and adds the keys to the given set
		err := dag.EnumerateChildren(ctx, verifyGetLinks, c, set.Visit)

		if err != nil {
			err = verboseCidError(err)
			return err
		}
	}

	return nil
}

// ColoredSet computes the set of nodes in the graph that are pinned by the
// pins in the given pinner.
func ColoredSet(ctx context.Context, pn pin.Pinner, ng ipld.NodeGetter, bestEffortRoots []cid.Cid, output chan<- Result) (*cid.Set, error) {
	// KeySet currently implemented in memory, in the future, may be bloom filter or
	// disk backed to conserve memory.
	errors := false
	gcs := cid.NewSet()
	getLinks := func(ctx context.Context, cid cid.Cid) ([]*ipld.Link, error) {
		links, err := ipld.GetLinks(ctx, ng, cid)
		if err != nil {
			errors = true
			select {
			case output <- Result{Error: &CannotFetchLinksError{cid, err}}:
			case <-ctx.Done():
			}
		}
		return links, nil
	}
	err := Descendants(ctx, getLinks, gcs, pn.RecursiveKeys())
	if err != nil {
		errors = true
		select {
		case output <- Result{Error: err}:
		case <-ctx.Done():
		}
	}

	bestEffortGetLinks := func(ctx context.Context, cid cid.Cid) ([]*ipld.Link, error) {
		links, err := ipld.GetLinks(ctx, ng, cid)
		if err != nil && err != ipld.ErrNotFound {
			errors = true
			select {
			case output <- Result{Error: &CannotFetchLinksError{cid, err}}:
			case <-ctx.Done():
			}
		}
		return links, nil
	}
	err = Descendants(ctx, bestEffortGetLinks, gcs, bestEffortRoots)
	if err != nil {
		errors = true
		select {
		case output <- Result{Error: err}:
		case <-ctx.Done():
		}
	}

	for _, k := range pn.DirectKeys() {
		gcs.Add(k)
	}

	err = Descendants(ctx, getLinks, gcs, pn.InternalPins())
	if err != nil {
		errors = true
		select {
		case output <- Result{Error: err}:
		case <-ctx.Done():
		}
	}

	if errors {
		return nil, ErrCannotFetchAllLinks
	}

	return gcs, nil
}

// ErrCannotFetchAllLinks is returned as the last Result in the GC output
// channel when there was a error creating the marked set because of a
// problem when finding descendants.
var ErrCannotFetchAllLinks = errors.New("garbage collection aborted: could not retrieve some links")

// ErrCannotDeleteSomeBlocks is returned when removing blocks marked for
// deletion fails as the last Result in GC output channel.
var ErrCannotDeleteSomeBlocks = errors.New("garbage collection incomplete: could not delete some blocks")

// CannotFetchLinksError provides detailed information about which links
// could not be fetched and can appear as a Result in the GC output channel.
type CannotFetchLinksError struct {
	Key cid.Cid
	Err error
}

// Error implements the error interface for this type with a useful
// message.
func (e *CannotFetchLinksError) Error() string {
	return fmt.Sprintf("could not retrieve links for %s: %s", e.Key, e.Err)
}

// CannotDeleteBlockError provides detailed information about which
// blocks could not be deleted and can appear as a Result in the GC output
// channel.
type CannotDeleteBlockError struct {
	Key cid.Cid
	Err error
}

// Error implements the error interface for this type with a
// useful message.
func (e *CannotDeleteBlockError) Error() string {
	return fmt.Sprintf("could not remove %s: %s", e.Key, e.Err)
}
<|endoftext|>


<commit_before>package test

import (
	"database/sql"
	"github.com/bwmarrin/discordgo"
	"github.com/gorilla/mux"
	"github.com/innovandalism/shodan"
)

type MockShodan struct {
	FuncGetDiscord      func() *discordgo.Session
	FuncGetMux          func() *mux.Router
	FuncGetDatabase     func() *sql.DB
	FuncGetCommandStack func() *shodan.CommandStack
	FuncGetRedis        func() shodan.KVS
}

func (ms *MockShodan) GetDiscord() *discordgo.Session {
	return ms.FuncGetDiscord()
}

func (ms *MockShodan) GetDatabase() *sql.DB {
	return ms.FuncGetDatabase()
}

func (ms *MockShodan) GetMux() *mux.Router {
	return ms.FuncGetMux()
}

func (ms *MockShodan) GetCommandStack() *shodan.CommandStack {
	return ms.FuncGetCommandStack()
}

func (ms *MockShodan) GetRedis() shodan.KVS {
	return ms.FuncGetRedis()
}
<commit_msg>Remove mocking tools<commit_after><|endoftext|>