<commit_before>/*
Copyright 2018 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 fakeclient

import (
	"context"
	"testing"
	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes/fake"
	"k8s.io/client-go/tools/cache"
)

// TestFakeClient demonstrates how to use a fake client with SharedInformerFactory in tests.
func TestFakeClient(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Create the fake client.
	client := fake.NewSimpleClientset()

	// We will create an informer that writes added pods to a channel.
	pods := make(chan *v1.Pod, 1)
	informers := informers.NewSharedInformerFactory(client, 0)
	podInformer := informers.Core().V1().Pods().Informer()
	podInformer.AddEventHandler(&cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			pod := obj.(*v1.Pod)
			t.Logf("pod added: %s/%s", pod.Namespace, pod.Name)
			pods <- pod
		},
	})

	// Make sure informers are running.
	informers.Start(ctx.Done())

	// This is not required in tests, but it serves as a proof-of-concept by
	// ensuring that the informer goroutine have warmed up and called List before
	// we send any events to it.
	cache.WaitForCacheSync(ctx.Done(), podInformer.HasSynced)

	// Inject an event into the fake client.
	p := &v1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "my-pod"}}
	_, err := client.CoreV1().Pods("test-ns").Create(context.TODO(), p, metav1.CreateOptions{})
	if err != nil {
		t.Fatalf("error injecting pod add: %v", err)
	}

	select {
	case pod := <-pods:
		t.Logf("Got pod from channel: %s/%s", pod.Namespace, pod.Name)
	case <-time.After(wait.ForeverTestTimeout):
		t.Error("Informer did not get the added pod")
	}
}
<commit_msg>fix the fake client example: how to handle a race between the fake client and informer<commit_after>/*
Copyright 2018 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 fakeclient

import (
	"context"
	"testing"
	"time"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes/fake"
	clienttesting "k8s.io/client-go/testing"
	"k8s.io/client-go/tools/cache"
)

// TestFakeClient demonstrates how to use a fake client with SharedInformerFactory in tests.
func TestFakeClient(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	watcherStarted := make(chan struct{})
	// Create the fake client.
	client := fake.NewSimpleClientset()
	// A catch-all watch reactor that allows us to inject the watcherStarted channel.
	client.PrependWatchReactor("*", func(action clienttesting.Action) (handled bool, ret watch.Interface, err error) {
		gvr := action.GetResource()
		ns := action.GetNamespace()
		watch, err := client.Tracker().Watch(gvr, ns)
		if err != nil {
			return false, nil, err
		}
		close(watcherStarted)
		return true, watch, nil
	})

	// We will create an informer that writes added pods to a channel.
	pods := make(chan *v1.Pod, 1)
	informers := informers.NewSharedInformerFactory(client, 0)
	podInformer := informers.Core().V1().Pods().Informer()
	podInformer.AddEventHandler(&cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			pod := obj.(*v1.Pod)
			t.Logf("pod added: %s/%s", pod.Namespace, pod.Name)
			pods <- pod
		},
	})

	// Make sure informers are running.
	informers.Start(ctx.Done())

	// This is not required in tests, but it serves as a proof-of-concept by
	// ensuring that the informer goroutine have warmed up and called List before
	// we send any events to it.
	cache.WaitForCacheSync(ctx.Done(), podInformer.HasSynced)

	// The fake client doesn't support resource version. Any writes to the client
	// after the informer's initial LIST and before the informer establishing the
	// watcher will be missed by the informer. Therefore we wait until the watcher
	// starts.
	// Note that the fake client isn't designed to work with informer. It
	// doesn't support resource version. It's encouraged to use a real client
	// in an integration/E2E test if you need to test complex behavior with
	// informer/controllers.
	<-watcherStarted
	// Inject an event into the fake client.
	p := &v1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "my-pod"}}
	_, err := client.CoreV1().Pods("test-ns").Create(context.TODO(), p, metav1.CreateOptions{})
	if err != nil {
		t.Fatalf("error injecting pod add: %v", err)
	}

	select {
	case pod := <-pods:
		t.Logf("Got pod from channel: %s/%s", pod.Namespace, pod.Name)
	case <-time.After(wait.ForeverTestTimeout):
		t.Error("Informer did not get the added pod")
	}
}
<|endoftext|>


<commit_before>package netutils

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
)

func NewHttpResponse(request *http.Request, statusCode int, body []byte, contentType string) *http.Response {
	resp := &http.Response{
		Status:     fmt.Sprintf("%d %s", statusCode, http.StatusText(statusCode)),
		StatusCode: statusCode,
		Proto:      "HTTP/1.0",
		ProtoMajor: 1,
		ProtoMinor: 0,
		Header:     make(http.Header),
	}
	resp.Header.Add("Content-Type", contentType)
	resp.Body = ioutil.NopCloser(bytes.NewBuffer(body))
	resp.ContentLength = int64(len(body))
	resp.Request = request
	return resp
}

func NewTextResponse(request *http.Request, statusCode int, body string) *http.Response {
	return NewHttpResponse(request, statusCode, []byte(body), "text/plain")
}
<commit_msg>Fix response<commit_after>package netutils

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
)

func NewHttpResponse(request *http.Request, statusCode int, body []byte, contentType string) *http.Response {
	resp := &http.Response{
		Status:     fmt.Sprintf("%d %s", statusCode, http.StatusText(statusCode)),
		StatusCode: statusCode,
		Proto:      "HTTP/1.0",
		ProtoMajor: 1,
		ProtoMinor: 0,
		Header:     make(http.Header),
	}
	resp.Header.Add("Content-Type", contentType)
	resp.Body = ioutil.NopCloser(bytes.NewBuffer(body))
	resp.ContentLength = int64(len(body))
	resp.Request = request
	return resp
}

func NewTextResponse(request *http.Request, statusCode int, body string) *http.Response {
	return NewHttpResponse(request, statusCode, []byte(body), "text/plain")
}

func NewJsonResponse(request *http.Request, statusCode int, message interface{}) *http.Response {
	bytes, err := json.Marshal(message)
	if err != nil {
		bytes = []byte("{}")
	}
	return NewHttpResponse(request, statusCode, bytes, "application/json")
}
<|endoftext|>


<commit_before>// Package cache provides MemCache as a TTL based simple cache in memory.
package cache

import (
	"github.com/hoveychen/go-utils"
	"sync"
	"time"
)

type MemCache struct {
	Items  map[string]*cachedItem
	lock   sync.RWMutex
	ticker *time.Ticker
}

type cachedItem struct {
	Payload    interface{}
	ExpireTime time.Time
}

func NewMemCache() *MemCache {
	c := &MemCache{
		Items: map[string]*cachedItem{},
	}
	go func() {
		ticker := time.NewTicker(time.Minute)
		for range ticker.C {
			c.checkExpire()
		}
	}()
	return c
}

func (c *MemCache) checkExpire() {
	c.lock.Lock()
	defer c.lock.Unlock()
	// TODO(Yuheng): Optimize the expire checking algorithm.
	now := goutils.GetNow()
	for key, item := range c.Items {
		if item.ExpireTime.Before(now) {
			delete(c.Items, key)
		}
	}
}

func (c *MemCache) Get(key string) interface{} {
	c.lock.RLock()
	defer c.lock.RUnlock()
	i, hit := c.Items[key]
	if hit {
		return i.Payload
	} else {
		return nil
	}
}

func (c *MemCache) UpsertWithTTL(key string, val interface{}, ttl time.Duration) {
	c.lock.Lock()
	defer c.lock.Unlock()
	i, hit := c.Items[key]
	if hit {
		i.Payload = val
		i.ExpireTime = goutils.GetNow().Add(ttl)
	} else {
		c.Items[key] = &cachedItem{
			Payload:    val,
			ExpireTime: goutils.GetNow().Add(ttl),
		}
	}
}

func (c *MemCache) Stop() {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.ticker.Stop()
	c.Items = nil
}
<commit_msg>[Memcache] 1. Change expiring calculation from start of inserting to the last access time. 2. Change API interface. NewMemcache accept from no parameter to one parameter "checkInterval"<commit_after>// Package cache provides MemCache as a TTL based simple cache in memory.
package cache

import (
	"sync"
	"time"

	"github.com/hoveychen/go-utils"
)

type MemCache struct {
	Items  map[string]*cachedItem
	lock   sync.RWMutex
	ticker *time.Ticker
}

type cachedItem struct {
	Payload    interface{}
	ExpireTime time.Time
	TTL        time.Duration
}

func NewMemCache(checkInterval time.Duration) *MemCache {
	c := &MemCache{
		Items: map[string]*cachedItem{},
	}
	go func() {
		ticker := time.NewTicker(checkInterval)
		for range ticker.C {
			c.checkExpire()
		}
	}()
	return c
}

func (c *MemCache) checkExpire() {
	c.lock.Lock()
	defer c.lock.Unlock()
	// TODO(Yuheng): Optimize the expire checking algorithm.
	now := goutils.GetNow()
	for key, item := range c.Items {
		if item.ExpireTime.Before(now) {
			delete(c.Items, key)
		}
	}
}

func (c *MemCache) Get(key string) interface{} {
	c.lock.RLock()
	defer c.lock.RUnlock()
	i, hit := c.Items[key]
	if hit {
		// It's hit. Extends the expiring time by another TTL.
		i.ExpireTime = goutils.GetNow().Add(i.TTL)
		return i.Payload
	} else {
		return nil
	}
}

func (c *MemCache) UpsertWithTTL(key string, val interface{}, ttl time.Duration) {
	c.lock.Lock()
	defer c.lock.Unlock()
	i, hit := c.Items[key]
	if hit {
		i.Payload = val
		i.TTL = ttl
		i.ExpireTime = goutils.GetNow().Add(ttl)
	} else {
		c.Items[key] = &cachedItem{
			Payload:    val,
			ExpireTime: goutils.GetNow().Add(ttl),
			TTL:        ttl,
		}
	}
}

func (c *MemCache) Stop() {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.ticker.Stop()
	c.Items = nil
}
<|endoftext|>


<commit_before>package captain // import "github.com/harbur/captain/captain"

import (
	"fmt"
	"os"

	"github.com/fsouza/go-dockerclient"
)

var endpoint = "unix:///var/run/docker.sock"
var client, _ = docker.NewClient(endpoint)

func buildImage(dockerfile string, image string, tag string) error {
	info("Building image %s:%s", image, tag)

	opts := docker.BuildImageOptions{
		Name:                image + ":" + tag,
		NoCache:             false,
		SuppressOutput:      false,
		RmTmpContainer:      true,
		ForceRmTmpContainer: true,
		OutputStream:        os.Stdout,
		ContextDir:          ".",
	}
	err := client.BuildImage(opts)
	if err != nil {
		fmt.Printf("%s", err)
	}
	return err
}

func tagImage(repo string, origin string, tag string) error {
	info("Tagging image %s:%s as %s:%s", repo, origin, repo, tag)
	// var imageID = getImageID(repo, origin)
	opts := docker.TagImageOptions{Repo: repo, Tag: tag, Force: true}
	err := client.TagImage(repo, opts)
	if err != nil {
		fmt.Printf("%s", err)
	}
	return err
}

func getImageID(repo string, tag string) string {
	images, _ := client.ListImages(docker.ListImagesOptions{})
	for _, image := range images {
		for _, b := range image.RepoTags {
			if b == repo+":"+tag {
				return image.ID
			}
		}
	}
	return ""
}

func imageExist(repo string, tag string) bool {
	images, _ := client.ListImages(docker.ListImagesOptions{})
	for _, image := range images {
		for _, b := range image.RepoTags {
			if b == repo+":"+tag {
				return true
			}
		}
	}
	return false
}
<commit_msg>Added Dockerfile variable on docker build command<commit_after>package captain // import "github.com/harbur/captain/captain"

import (
	"fmt"
	"os"

	"github.com/fsouza/go-dockerclient"
)

var endpoint = "unix:///var/run/docker.sock"
var client, _ = docker.NewClient(endpoint)

func buildImage(dockerfile string, image string, tag string) error {
	info("Building image %s:%s", image, tag)

	opts := docker.BuildImageOptions{
		Name:                image + ":" + tag,
		Dockerfile:          dockerfile,
		NoCache:             false,
		SuppressOutput:      false,
		RmTmpContainer:      true,
		ForceRmTmpContainer: true,
		OutputStream:        os.Stdout,
		ContextDir:          ".",
	}
	err := client.BuildImage(opts)
	if err != nil {
		fmt.Printf("%s", err)
	}
	return err
}

func tagImage(repo string, origin string, tag string) error {
	info("Tagging image %s:%s as %s:%s", repo, origin, repo, tag)
	// var imageID = getImageID(repo, origin)
	opts := docker.TagImageOptions{Repo: repo, Tag: tag, Force: true}
	err := client.TagImage(repo, opts)
	if err != nil {
		fmt.Printf("%s", err)
	}
	return err
}

func getImageID(repo string, tag string) string {
	images, _ := client.ListImages(docker.ListImagesOptions{})
	for _, image := range images {
		for _, b := range image.RepoTags {
			if b == repo+":"+tag {
				return image.ID
			}
		}
	}
	return ""
}

func imageExist(repo string, tag string) bool {
	images, _ := client.ListImages(docker.ListImagesOptions{})
	for _, image := range images {
		for _, b := range image.RepoTags {
			if b == repo+":"+tag {
				return true
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>package readpassword

import (
	"os"
	"os/exec"
	"testing"
)

func TestPassfile(t *testing.T) {
	testcases := []struct {
		file string
		want string
	}{
		{"mypassword.txt", "mypassword"},
		{"mypassword_garbage.txt", "mypassword"},
		{"mypassword_missing_newline.txt", "mypassword"},
	}
	for _, tc := range testcases {
		pw := readPassFile("passfile_test_files/" + tc.file)
		if string(pw) != tc.want {
			t.Errorf("Wrong result: want=%q have=%q", tc.want, pw)
		}
	}
}

// readPassFile() should exit instead of returning an empty string.
//
// The TEST_SLAVE magic is explained at
// https://talks.golang.org/2014/testing.slide#23 .
func TestPassfileEmpty(t *testing.T) {
	if os.Getenv("TEST_SLAVE") == "1" {
		readPassFile("passfile_test_files/empty.txt")
		return
	}
	cmd := exec.Command(os.Args[0], "-test.run=TestPassfileEmpty$")
	cmd.Env = append(os.Environ(), "TEST_SLAVE=1")
	err := cmd.Run()
	if err != nil {
		return
	}
	t.Fatal("should have exited")
}

// File containing just a newline.
// readPassFile() should exit instead of returning an empty string.
//
// The TEST_SLAVE magic is explained at
// https://talks.golang.org/2014/testing.slide#23 .
func TestPassfileNewline(t *testing.T) {
	if os.Getenv("TEST_SLAVE") == "1" {
		readPassFile("passfile_test_files/newline.txt")
		return
	}
	cmd := exec.Command(os.Args[0], "-test.run=TestPassfileEmpty$")
	cmd.Env = append(os.Environ(), "TEST_SLAVE=1")
	err := cmd.Run()
	if err != nil {
		return
	}
	t.Fatal("should have exited")
}

// File containing "\ngarbage".
// readPassFile() should exit instead of returning an empty string.
//
// The TEST_SLAVE magic is explained at
// https://talks.golang.org/2014/testing.slide#23 .
func TestPassfileEmptyFirstLine(t *testing.T) {
	if os.Getenv("TEST_SLAVE") == "1" {
		readPassFile("passfile_test_files/empty_first_line.txt")
		return
	}
	cmd := exec.Command(os.Args[0], "-test.run=TestPassfileEmptyFirstLine$")
	cmd.Env = append(os.Environ(), "TEST_SLAVE=1")
	err := cmd.Run()
	if err != nil {
		return
	}
	t.Fatal("should have exited")
}
<commit_msg>tests: fix TestPassfileNewline<commit_after>package readpassword

import (
	"os"
	"os/exec"
	"testing"
)

func TestPassfile(t *testing.T) {
	testcases := []struct {
		file string
		want string
	}{
		{"mypassword.txt", "mypassword"},
		{"mypassword_garbage.txt", "mypassword"},
		{"mypassword_missing_newline.txt", "mypassword"},
	}
	for _, tc := range testcases {
		pw := readPassFile("passfile_test_files/" + tc.file)
		if string(pw) != tc.want {
			t.Errorf("Wrong result: want=%q have=%q", tc.want, pw)
		}
	}
}

// readPassFile() should exit instead of returning an empty string.
//
// The TEST_SLAVE magic is explained at
// https://talks.golang.org/2014/testing.slide#23 .
func TestPassfileEmpty(t *testing.T) {
	if os.Getenv("TEST_SLAVE") == "1" {
		readPassFile("passfile_test_files/empty.txt")
		return
	}
	cmd := exec.Command(os.Args[0], "-test.run=TestPassfileEmpty$")
	cmd.Env = append(os.Environ(), "TEST_SLAVE=1")
	err := cmd.Run()
	if err != nil {
		return
	}
	t.Fatal("should have exited")
}

// File containing just a newline.
// readPassFile() should exit instead of returning an empty string.
//
// The TEST_SLAVE magic is explained at
// https://talks.golang.org/2014/testing.slide#23 .
func TestPassfileNewline(t *testing.T) {
	if os.Getenv("TEST_SLAVE") == "1" {
		readPassFile("passfile_test_files/newline.txt")
		return
	}
	cmd := exec.Command(os.Args[0], "-test.run=TestPassfileNewline$")
	cmd.Env = append(os.Environ(), "TEST_SLAVE=1")
	err := cmd.Run()
	if err != nil {
		return
	}
	t.Fatal("should have exited")
}

// File containing "\ngarbage".
// readPassFile() should exit instead of returning an empty string.
//
// The TEST_SLAVE magic is explained at
// https://talks.golang.org/2014/testing.slide#23 .
func TestPassfileEmptyFirstLine(t *testing.T) {
	if os.Getenv("TEST_SLAVE") == "1" {
		readPassFile("passfile_test_files/empty_first_line.txt")
		return
	}
	cmd := exec.Command(os.Args[0], "-test.run=TestPassfileEmptyFirstLine$")
	cmd.Env = append(os.Environ(), "TEST_SLAVE=1")
	err := cmd.Run()
	if err != nil {
		return
	}
	t.Fatal("should have exited")
}
<|endoftext|>


<commit_before>package sessiontest

import (
	"encoding/json"
	"reflect"
	"testing"

	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/test"
	"github.com/privacybydesign/irmago/irmaclient"
	"github.com/privacybydesign/irmago/server"
	"github.com/stretchr/testify/require"
)

type sessionOption int

const (
	sessionOptionUpdatedIrmaConfiguration = iota
	sessionOptionUnsatisfiableRequest
)

type requestorSessionResult struct {
	*server.SessionResult
	Missing irmaclient.MissingAttributes
}

func requestorSessionHelper(t *testing.T, request irma.SessionRequest, client *irmaclient.Client, options ...sessionOption) *requestorSessionResult {
	if client == nil {
		client, _ = parseStorage(t)
		defer test.ClearTestStorage(t)
	}

	StartIrmaServer(t, len(options) == 1 && options[0] == sessionOptionUpdatedIrmaConfiguration)
	defer StopIrmaServer()

	clientChan := make(chan *SessionResult)
	serverChan := make(chan *server.SessionResult)

	qr, token, err := irmaServer.StartSession(request, func(result *server.SessionResult) {
		serverChan <- result
	})
	require.NoError(t, err)

	var h irmaclient.Handler
	if len(options) == 1 && options[0] == sessionOptionUnsatisfiableRequest {
		h = UnsatisfiableTestHandler{TestHandler{t, clientChan, client, nil}}
	} else {
		h = TestHandler{t, clientChan, client, nil}
	}
	j, err := json.Marshal(qr)
	require.NoError(t, err)
	client.NewSession(string(j), h)
	clientResult := <-clientChan
	if clientResult != nil {
		require.NoError(t, clientResult.Err)
	}

	if len(options) == 1 && options[0] == sessionOptionUnsatisfiableRequest {
		return &requestorSessionResult{nil, clientResult.Missing}
	} else {
		serverResult := <-serverChan
		require.Equal(t, token, serverResult.Token)
		return &requestorSessionResult{serverResult, nil}
	}
}

// Check that nonexistent IRMA identifiers in the session request fail the session
func TestRequestorInvalidRequest(t *testing.T) {
	StartIrmaServer(t, false)
	defer StopIrmaServer()
	_, _, err := irmaServer.StartSession(irma.NewDisclosureRequest(
		irma.NewAttributeTypeIdentifier("irma-demo.RU.foo.bar"),
		irma.NewAttributeTypeIdentifier("irma-demo.baz.qux.abc"),
	), nil)
	require.Error(t, err)
}

func TestRequestorSignatureSession(t *testing.T) {
	client, _ := parseStorage(t)
	id := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")
	serverResult := requestorSessionHelper(t, irma.NewSignatureRequest("message", id), client)

	require.Nil(t, serverResult.Err)
	require.Equal(t, irma.ProofStatusValid, serverResult.ProofStatus)
	require.NotEmpty(t, serverResult.Disclosed)
	require.Equal(t, id, serverResult.Disclosed[0][0].Identifier)
	require.Equal(t, "456", serverResult.Disclosed[0][0].Value["en"])

	// Load the updated scheme in which an attribute was added to the studentCard credential type
	schemeid := irma.NewSchemeManagerIdentifier("irma-demo")
	client.Configuration.SchemeManagers[schemeid].URL = "http://localhost:48681/irma_configuration_updated/irma-demo"
	require.NoError(t, client.Configuration.UpdateSchemeManager(schemeid, nil))
	require.NoError(t, client.Configuration.ParseFolder())
	require.Contains(t, client.Configuration.AttributeTypes, irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.newAttribute"))

	// Check that the just created credential is still valid after the new attribute has been added
	_, status, err := serverResult.Signature.Verify(client.Configuration, nil)
	require.NoError(t, err)
	require.Equal(t, irma.ProofStatusValid, status)
}

func TestRequestorDisclosureSession(t *testing.T) {
	id := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")
	request := irma.NewDisclosureRequest(id)
	serverResult := testRequestorDisclosure(t, request)
	require.Len(t, serverResult.Disclosed, 1)
	require.Equal(t, id, serverResult.Disclosed[0][0].Identifier)
	require.Equal(t, "456", serverResult.Disclosed[0][0].Value["en"])
}

func TestRequestorDisclosureMultipleAttrs(t *testing.T) {
	request := irma.NewDisclosureRequest(
		irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID"),
		irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.level"),
	)
	serverResult := testRequestorDisclosure(t, request)
	require.Len(t, serverResult.Disclosed, 2)
}

func testRequestorDisclosure(t *testing.T, request *irma.DisclosureRequest) *server.SessionResult {
	serverResult := requestorSessionHelper(t, request, nil)
	require.Nil(t, serverResult.Err)
	require.Equal(t, irma.ProofStatusValid, serverResult.ProofStatus)
	return serverResult.SessionResult
}

func TestRequestorIssuanceSession(t *testing.T) {
	testRequestorIssuance(t, false)
}

func TestRequestorCombinedSessionMultipleAttributes(t *testing.T) {
	var ir irma.IssuanceRequest
	require.NoError(t, irma.UnmarshalValidate([]byte(`{
		"type":"issuing",
		"credentials": [
			{
				"credential":"irma-demo.MijnOverheid.root",
				"attributes" : {
					"BSN":"12345"
				}
			}
		],
		"disclose" : [
			{
				"label":"Initialen",
				"attributes":["irma-demo.RU.studentCard.studentCardNumber"]
			},
			{
				"label":"Achternaam",
				"attributes" : ["irma-demo.RU.studentCard.studentID"]
			},
			{
				"label":"Geboortedatum",
				"attributes":["irma-demo.RU.studentCard.university"]
			}
		]
	}`), &ir))

	require.Equal(t, server.StatusDone, requestorSessionHelper(t, &ir, nil).Status)
}

func testRequestorIssuance(t *testing.T, keyshare bool) {
	attrid := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")
	request := irma.NewIssuanceRequest([]*irma.CredentialRequest{{
		CredentialTypeID: irma.NewCredentialTypeIdentifier("irma-demo.RU.studentCard"),
		Attributes: map[string]string{
			"university":        "Radboud",
			"studentCardNumber": "31415927",
			"studentID":         "s1234567",
			"level":             "42",
		},
	}, {
		CredentialTypeID: irma.NewCredentialTypeIdentifier("irma-demo.MijnOverheid.root"),
		Attributes: map[string]string{
			"BSN": "299792458",
		},
	}}, attrid)
	if keyshare {
		request.Credentials = append(request.Credentials, &irma.CredentialRequest{
			CredentialTypeID: irma.NewCredentialTypeIdentifier("test.test.mijnirma"),
			Attributes:       map[string]string{"email": "testusername"},
		})
	}

	result := requestorSessionHelper(t, request, nil)
	require.Nil(t, result.Err)
	require.Equal(t, irma.ProofStatusValid, result.ProofStatus)
	require.NotEmpty(t, result.Disclosed)
	require.Equal(t, attrid, result.Disclosed[0][0].Identifier)
	require.Equal(t, "456", result.Disclosed[0][0].Value["en"])
}

func TestConDisCon(t *testing.T) {
	client, _ := parseStorage(t)
	ir := getMultipleIssuanceRequest()
	ir.Credentials = append(ir.Credentials, &irma.CredentialRequest{
		Validity:         ir.Credentials[0].Validity,
		CredentialTypeID: irma.NewCredentialTypeIdentifier("irma-demo.MijnOverheid.fullName"),
		Attributes: map[string]string{
			"firstnames": "Jan Hendrik",
			"firstname":  "Jan",
			"familyname": "Klaassen",
			"prefix":     "van",
		},
	})
	requestorSessionHelper(t, ir, client)

	dr := irma.NewDisclosureRequest()
	dr.Disclose = irma.AttributeConDisCon{
		irma.AttributeDisCon{
			irma.AttributeCon{
				irma.NewAttributeRequest("irma-demo.MijnOverheid.root.BSN"),
				irma.NewAttributeRequest("irma-demo.MijnOverheid.fullName.firstname"),
				irma.NewAttributeRequest("irma-demo.MijnOverheid.fullName.familyname"),
			},
			irma.AttributeCon{
				irma.NewAttributeRequest("irma-demo.RU.studentCard.studentID"),
				irma.NewAttributeRequest("irma-demo.RU.studentCard.university"),
			},
		},
		//irma.AttributeDisCon{
		//	irma.AttributeCon{
		//		irma.NewAttributeRequest("irma-demo.MijnOverheid.fullName.firstname"),
		//		irma.NewAttributeRequest("irma-demo.MijnOverheid.fullName.familyname"),
		//	},
		//},
	}

	requestorSessionHelper(t, dr, client)
}

func TestOptionalDisclosure(t *testing.T) {
	client, _ := parseStorage(t)
	university := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.university")
	studentid := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")

	radboud := "Radboud"
	attrs1 := irma.AttributeConDisCon{
		irma.AttributeDisCon{ // Including one non-optional disjunction is required in disclosure and signature sessions
			irma.AttributeCon{irma.AttributeRequest{Type: university}},
		},
		irma.AttributeDisCon{
			irma.AttributeCon{},
			irma.AttributeCon{irma.AttributeRequest{Type: studentid}},
		},
	}
	disclosed1 := [][]*irma.DisclosedAttribute{
		{
			{
				RawValue:   &radboud,
				Value:      map[string]string{"": radboud, "en": radboud, "nl": radboud},
				Identifier: irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.university"),
				Status:     irma.AttributeProofStatusPresent,
			},
		},
		{},
	}
	attrs2 := irma.AttributeConDisCon{ // In issuance sessions, it is allowed that all disjunctions are optional
		irma.AttributeDisCon{
			irma.AttributeCon{},
			irma.AttributeCon{irma.AttributeRequest{Type: studentid}},
		},
	}
	disclosed2 := [][]*irma.DisclosedAttribute{{}}

	tests := []struct {
		request   irma.SessionRequest
		attrs     irma.AttributeConDisCon
		disclosed [][]*irma.DisclosedAttribute
	}{
		{irma.NewDisclosureRequest(), attrs1, disclosed1},
		{irma.NewSignatureRequest("message"), attrs1, disclosed1},
		{getIssuanceRequest(true), attrs1, disclosed1},
		{getIssuanceRequest(true), attrs2, disclosed2},
	}

	for _, args := range tests {
		args.request.Disclosure().Disclose = args.attrs

		// TestHandler always prefers the first option when given any choice, so it will not disclose the optional attribute
		result := requestorSessionHelper(t, args.request, client)
		require.True(t, reflect.DeepEqual(args.disclosed, result.Disclosed))
	}
}
<commit_msg>test: have TestOptionalDisclosure also compare issuance time (cf. 82a4ad4)<commit_after>package sessiontest

import (
	"encoding/json"
	"reflect"
	"testing"

	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/test"
	"github.com/privacybydesign/irmago/irmaclient"
	"github.com/privacybydesign/irmago/server"
	"github.com/stretchr/testify/require"
)

type sessionOption int

const (
	sessionOptionUpdatedIrmaConfiguration = iota
	sessionOptionUnsatisfiableRequest
)

type requestorSessionResult struct {
	*server.SessionResult
	Missing irmaclient.MissingAttributes
}

func requestorSessionHelper(t *testing.T, request irma.SessionRequest, client *irmaclient.Client, options ...sessionOption) *requestorSessionResult {
	if client == nil {
		client, _ = parseStorage(t)
		defer test.ClearTestStorage(t)
	}

	StartIrmaServer(t, len(options) == 1 && options[0] == sessionOptionUpdatedIrmaConfiguration)
	defer StopIrmaServer()

	clientChan := make(chan *SessionResult)
	serverChan := make(chan *server.SessionResult)

	qr, token, err := irmaServer.StartSession(request, func(result *server.SessionResult) {
		serverChan <- result
	})
	require.NoError(t, err)

	var h irmaclient.Handler
	if len(options) == 1 && options[0] == sessionOptionUnsatisfiableRequest {
		h = UnsatisfiableTestHandler{TestHandler{t, clientChan, client, nil}}
	} else {
		h = TestHandler{t, clientChan, client, nil}
	}
	j, err := json.Marshal(qr)
	require.NoError(t, err)
	client.NewSession(string(j), h)
	clientResult := <-clientChan
	if clientResult != nil {
		require.NoError(t, clientResult.Err)
	}

	if len(options) == 1 && options[0] == sessionOptionUnsatisfiableRequest {
		return &requestorSessionResult{nil, clientResult.Missing}
	} else {
		serverResult := <-serverChan
		require.Equal(t, token, serverResult.Token)
		return &requestorSessionResult{serverResult, nil}
	}
}

// Check that nonexistent IRMA identifiers in the session request fail the session
func TestRequestorInvalidRequest(t *testing.T) {
	StartIrmaServer(t, false)
	defer StopIrmaServer()
	_, _, err := irmaServer.StartSession(irma.NewDisclosureRequest(
		irma.NewAttributeTypeIdentifier("irma-demo.RU.foo.bar"),
		irma.NewAttributeTypeIdentifier("irma-demo.baz.qux.abc"),
	), nil)
	require.Error(t, err)
}

func TestRequestorSignatureSession(t *testing.T) {
	client, _ := parseStorage(t)
	id := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")
	serverResult := requestorSessionHelper(t, irma.NewSignatureRequest("message", id), client)

	require.Nil(t, serverResult.Err)
	require.Equal(t, irma.ProofStatusValid, serverResult.ProofStatus)
	require.NotEmpty(t, serverResult.Disclosed)
	require.Equal(t, id, serverResult.Disclosed[0][0].Identifier)
	require.Equal(t, "456", serverResult.Disclosed[0][0].Value["en"])

	// Load the updated scheme in which an attribute was added to the studentCard credential type
	schemeid := irma.NewSchemeManagerIdentifier("irma-demo")
	client.Configuration.SchemeManagers[schemeid].URL = "http://localhost:48681/irma_configuration_updated/irma-demo"
	require.NoError(t, client.Configuration.UpdateSchemeManager(schemeid, nil))
	require.NoError(t, client.Configuration.ParseFolder())
	require.Contains(t, client.Configuration.AttributeTypes, irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.newAttribute"))

	// Check that the just created credential is still valid after the new attribute has been added
	_, status, err := serverResult.Signature.Verify(client.Configuration, nil)
	require.NoError(t, err)
	require.Equal(t, irma.ProofStatusValid, status)
}

func TestRequestorDisclosureSession(t *testing.T) {
	id := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")
	request := irma.NewDisclosureRequest(id)
	serverResult := testRequestorDisclosure(t, request)
	require.Len(t, serverResult.Disclosed, 1)
	require.Equal(t, id, serverResult.Disclosed[0][0].Identifier)
	require.Equal(t, "456", serverResult.Disclosed[0][0].Value["en"])
}

func TestRequestorDisclosureMultipleAttrs(t *testing.T) {
	request := irma.NewDisclosureRequest(
		irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID"),
		irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.level"),
	)
	serverResult := testRequestorDisclosure(t, request)
	require.Len(t, serverResult.Disclosed, 2)
}

func testRequestorDisclosure(t *testing.T, request *irma.DisclosureRequest) *server.SessionResult {
	serverResult := requestorSessionHelper(t, request, nil)
	require.Nil(t, serverResult.Err)
	require.Equal(t, irma.ProofStatusValid, serverResult.ProofStatus)
	return serverResult.SessionResult
}

func TestRequestorIssuanceSession(t *testing.T) {
	testRequestorIssuance(t, false)
}

func TestRequestorCombinedSessionMultipleAttributes(t *testing.T) {
	var ir irma.IssuanceRequest
	require.NoError(t, irma.UnmarshalValidate([]byte(`{
		"type":"issuing",
		"credentials": [
			{
				"credential":"irma-demo.MijnOverheid.root",
				"attributes" : {
					"BSN":"12345"
				}
			}
		],
		"disclose" : [
			{
				"label":"Initialen",
				"attributes":["irma-demo.RU.studentCard.studentCardNumber"]
			},
			{
				"label":"Achternaam",
				"attributes" : ["irma-demo.RU.studentCard.studentID"]
			},
			{
				"label":"Geboortedatum",
				"attributes":["irma-demo.RU.studentCard.university"]
			}
		]
	}`), &ir))

	require.Equal(t, server.StatusDone, requestorSessionHelper(t, &ir, nil).Status)
}

func testRequestorIssuance(t *testing.T, keyshare bool) {
	attrid := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")
	request := irma.NewIssuanceRequest([]*irma.CredentialRequest{{
		CredentialTypeID: irma.NewCredentialTypeIdentifier("irma-demo.RU.studentCard"),
		Attributes: map[string]string{
			"university":        "Radboud",
			"studentCardNumber": "31415927",
			"studentID":         "s1234567",
			"level":             "42",
		},
	}, {
		CredentialTypeID: irma.NewCredentialTypeIdentifier("irma-demo.MijnOverheid.root"),
		Attributes: map[string]string{
			"BSN": "299792458",
		},
	}}, attrid)
	if keyshare {
		request.Credentials = append(request.Credentials, &irma.CredentialRequest{
			CredentialTypeID: irma.NewCredentialTypeIdentifier("test.test.mijnirma"),
			Attributes:       map[string]string{"email": "testusername"},
		})
	}

	result := requestorSessionHelper(t, request, nil)
	require.Nil(t, result.Err)
	require.Equal(t, irma.ProofStatusValid, result.ProofStatus)
	require.NotEmpty(t, result.Disclosed)
	require.Equal(t, attrid, result.Disclosed[0][0].Identifier)
	require.Equal(t, "456", result.Disclosed[0][0].Value["en"])
}

func TestConDisCon(t *testing.T) {
	client, _ := parseStorage(t)
	ir := getMultipleIssuanceRequest()
	ir.Credentials = append(ir.Credentials, &irma.CredentialRequest{
		Validity:         ir.Credentials[0].Validity,
		CredentialTypeID: irma.NewCredentialTypeIdentifier("irma-demo.MijnOverheid.fullName"),
		Attributes: map[string]string{
			"firstnames": "Jan Hendrik",
			"firstname":  "Jan",
			"familyname": "Klaassen",
			"prefix":     "van",
		},
	})
	requestorSessionHelper(t, ir, client)

	dr := irma.NewDisclosureRequest()
	dr.Disclose = irma.AttributeConDisCon{
		irma.AttributeDisCon{
			irma.AttributeCon{
				irma.NewAttributeRequest("irma-demo.MijnOverheid.root.BSN"),
				irma.NewAttributeRequest("irma-demo.MijnOverheid.fullName.firstname"),
				irma.NewAttributeRequest("irma-demo.MijnOverheid.fullName.familyname"),
			},
			irma.AttributeCon{
				irma.NewAttributeRequest("irma-demo.RU.studentCard.studentID"),
				irma.NewAttributeRequest("irma-demo.RU.studentCard.university"),
			},
		},
		//irma.AttributeDisCon{
		//	irma.AttributeCon{
		//		irma.NewAttributeRequest("irma-demo.MijnOverheid.fullName.firstname"),
		//		irma.NewAttributeRequest("irma-demo.MijnOverheid.fullName.familyname"),
		//	},
		//},
	}

	requestorSessionHelper(t, dr, client)
}

func TestOptionalDisclosure(t *testing.T) {
	client, _ := parseStorage(t)
	university := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.university")
	studentid := irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")

	radboud := "Radboud"
	attrs1 := irma.AttributeConDisCon{
		irma.AttributeDisCon{ // Including one non-optional disjunction is required in disclosure and signature sessions
			irma.AttributeCon{irma.AttributeRequest{Type: university}},
		},
		irma.AttributeDisCon{
			irma.AttributeCon{},
			irma.AttributeCon{irma.AttributeRequest{Type: studentid}},
		},
	}
	disclosed1 := [][]*irma.DisclosedAttribute{
		{
			{
				RawValue:     &radboud,
				Value:        map[string]string{"": radboud, "en": radboud, "nl": radboud},
				Identifier:   irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.university"),
				Status:       irma.AttributeProofStatusPresent,
				IssuanceTime: irma.Timestamp(client.Attributes(university.CredentialTypeIdentifier(), 0).SigningDate()),
			},
		},
		{},
	}
	attrs2 := irma.AttributeConDisCon{ // In issuance sessions, it is allowed that all disjunctions are optional
		irma.AttributeDisCon{
			irma.AttributeCon{},
			irma.AttributeCon{irma.AttributeRequest{Type: studentid}},
		},
	}
	disclosed2 := [][]*irma.DisclosedAttribute{{}}

	tests := []struct {
		request   irma.SessionRequest
		attrs     irma.AttributeConDisCon
		disclosed [][]*irma.DisclosedAttribute
	}{
		{irma.NewDisclosureRequest(), attrs1, disclosed1},
		{irma.NewSignatureRequest("message"), attrs1, disclosed1},
		{getIssuanceRequest(true), attrs1, disclosed1},
		{getIssuanceRequest(true), attrs2, disclosed2},
	}

	for _, args := range tests {
		args.request.Disclosure().Disclose = args.attrs

		// TestHandler always prefers the first option when given any choice, so it will not disclose the optional attribute
		result := requestorSessionHelper(t, args.request, client)
		require.True(t, reflect.DeepEqual(args.disclosed, result.Disclosed))
	}
}
<|endoftext|>


<commit_before>package issuecomment

import (
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"
	"sync"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/auth"
	"github.com/parkr/auto-reply/ctx"
	"github.com/parkr/changelog"
)

// changelogCategory is a changelog category, like "Site Enhancements" and such.
type changelogCategory struct {
	Prefix, Slug, Section string
	Labels                []string
}

var (
	mergeCommentRegexp = regexp.MustCompile("@[a-zA-Z-_]+: (merge|:shipit:|:ship:)( \\+([a-zA-Z-_ ]+))?")

	categories = []changelogCategory{
		changelogCategory{
			Prefix:  "major",
			Slug:    "major-enhancements",
			Section: "Major Enhancements",
			Labels:  []string{"feature"},
		},
		changelogCategory{
			Prefix:  "minor",
			Slug:    "minor-enhancements",
			Section: "Minor Enhancements",
			Labels:  []string{"enhancement"},
		},
		changelogCategory{
			Prefix:  "bug",
			Slug:    "bug-fixes",
			Section: "Bug Fixes",
			Labels:  []string{"bug", "fix"},
		},
		changelogCategory{
			Prefix:  "dev",
			Slug:    "development-fixes",
			Section: "Development Fixes",
			Labels:  []string{"internal", "fix"},
		},
		changelogCategory{
			Prefix:  "port",
			Slug:    "forward-ports",
			Section: "Forward Ports",
			Labels:  []string{"forward-port"},
		},
		changelogCategory{
			Prefix:  "site",
			Slug:    "site-enhancements",
			Section: "Site Enhancements",
			Labels:  []string{"documentation"},
		},
	}
)

func MergeAndLabel(context *ctx.Context, payload interface{}) error {
	event, ok := payload.(*github.IssueCommentEvent)
	if !ok {
		return context.NewError("MergeAndLabel: not an issue comment event")
	}

	// Is this a pull request?
	if event.Issue != nil && event.Issue.PullRequestLinks != nil {
		return context.NewError("MergeAndLabel: not a pull request")
	}

	var changeSectionLabel string
	isReq, labelFromComment := parseMergeRequestComment(*event.Comment.Body)

	// Is It a merge request comment?
	if !isReq {
		return context.NewError("MergeAndLabel: not a merge request comment")
	}

	if os.Getenv("AUTO_REPLY_DEBUG") == "true" {
		log.Println("MergeAndLabel: received event:", event)
	}

	var wg sync.WaitGroup

	owner := *event.Repo.Owner.Login
	repo := *event.Repo.Name
	number := *event.Issue.Number
	ref := fmt.Sprintf("%s/%s#%d", owner, repo, number)

	// Does the user have merge/label abilities?
	if !auth.CommenterHasPushAccess(context, *event) {
		log.Printf("%s isn't authenticated to merge anything on %s", *event.Comment.User.Login, *event.Repo.FullName)
		return errors.New("commenter isn't allowed to merge")
	}

	// Should it be labeled?
	if labelFromComment != "" {
		changeSectionLabel = sectionForLabel(labelFromComment)
	} else {
		changeSectionLabel = "none"
	}
	fmt.Printf("changeSectionLabel = '%s'\n", changeSectionLabel)

	// Merge
	commitMsg := fmt.Sprintf("Merge pull request %v", number)
	_, _, mergeErr := context.GitHub.PullRequests.Merge(owner, repo, number, commitMsg)
	if mergeErr != nil {
		return context.NewError("MergeAndLabel: error merging %s: %v", ref, mergeErr)
	}

	// Delete branch
	repoInfo, _, getRepoErr := context.GitHub.PullRequests.Get(owner, repo, number)
	if getRepoErr != nil {
		return context.NewError("MergeAndLabel: error getting PR info %s: %v", ref, getRepoErr)
	}

	// Delete branch
	if deletableRef(repoInfo, owner) {
		wg.Add(1)
		go func() {
			ref := fmt.Sprintf("heads/%s", *repoInfo.Head.Ref)
			_, deleteBranchErr := context.GitHub.Git.DeleteRef(owner, repo, ref)
			if deleteBranchErr != nil {
				fmt.Printf("MergeAndLabel: error deleting branch %v\n", mergeErr)
			}
			wg.Done()
		}()
	}

	wg.Add(1)
	go func() {
		err := addLabelsForSubsection(context, owner, repo, number, changeSectionLabel)
		if err != nil {
			fmt.Printf("MergeAndLabel: error applying labels: %v\n", err)
		}
		wg.Done()
	}()

	wg.Add(1)
	go func() {
		// Read History.markdown, add line to appropriate change section
		historyFileContents, historySHA := getHistoryContents(context, owner, repo)

		// Add merge reference to history
		newHistoryFileContents := addMergeReference(historyFileContents, changeSectionLabel, *repoInfo.Title, number)

		// Commit change to History.markdown
		commitErr := commitHistoryFile(context, historySHA, owner, repo, number, newHistoryFileContents)
		if commitErr != nil {
			fmt.Printf("comments: error committing updated history %v\n", mergeErr)
		}
		wg.Done()
	}()

	wg.Wait()

	return nil
}

func parseMergeRequestComment(commentBody string) (bool, string) {
	matches := mergeCommentRegexp.FindAllStringSubmatch(commentBody, -1)
	if matches == nil || matches[0] == nil {
		return false, ""
	}

	var label string
	if len(matches[0]) >= 4 {
		if labelFromComment := matches[0][3]; labelFromComment != "" {
			label = downcaseAndHyphenize(labelFromComment)
		}
	}

	return true, normalizeLabel(label)
}

func downcaseAndHyphenize(label string) string {
	return strings.Replace(strings.ToLower(label), " ", "-", -1)
}

func normalizeLabel(label string) string {
	for _, category := range categories {
		if strings.HasPrefix(label, category.Prefix) {
			return category.Slug
		}
	}

	return label
}

func sectionForLabel(slug string) string {
	for _, category := range categories {
		if slug == category.Slug {
			return category.Section
		}
	}

	return slug
}

func labelsForSubsection(changeSectionLabel string) []string {
	for _, category := range categories {
		if changeSectionLabel == category.Section {
			return category.Labels
		}
	}

	return []string{}
}

func selectSectionLabel(labels []github.Label) string {
	for _, label := range labels {
		if sectionForLabel(*label.Name) != *label.Name {
			return *label.Name
		}
	}
	return ""
}

func containsChangeLabel(commentBody string) bool {
	_, labelFromComment := parseMergeRequestComment(commentBody)
	return labelFromComment != ""
}

func addLabelsForSubsection(context *ctx.Context, owner, repo string, number int, changeSectionLabel string) error {
	labels := labelsForSubsection(changeSectionLabel)

	if len(labels) < 1 {
		return fmt.Errorf("no labels for changeSectionLabel='%s'", changeSectionLabel)
	}

	_, _, err := context.GitHub.Issues.AddLabelsToIssue(owner, repo, number, labels)
	return err
}

func getHistoryContents(context *ctx.Context, owner, repo string) (content, sha string) {
	contents, _, _, err := context.GitHub.Repositories.GetContents(
		owner,
		repo,
		"History.markdown",
		&github.RepositoryContentGetOptions{Ref: "heads/master"},
	)
	if err != nil {
		fmt.Printf("comments: error getting History.markdown %v\n", err)
		return "", ""
	}
	return base64Decode(*contents.Content), *contents.SHA
}

func base64Decode(encoded string) string {
	decoded, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		fmt.Printf("comments: error decoding string: %v\n", err)
		return ""
	}
	return string(decoded)
}

func addMergeReference(historyFileContents, changeSectionLabel, prTitle string, number int) string {
	changes, err := changelog.NewChangelogFromReader(strings.NewReader(historyFileContents))
	if historyFileContents == "" {
		err = nil
		changes = changelog.NewChangelog()
	}
	if err != nil {
		fmt.Printf("comments: error %v\n", err)
		return historyFileContents
	}

	changeLine := &changelog.ChangeLine{
		Summary:   prTitle,
		Reference: fmt.Sprintf("#%d", number),
	}

	// Put either directly in the version history or in a subsection.
	if changeSectionLabel == "none" {
		changes.AddLineToVersion("HEAD", changeLine)
	} else {
		changes.AddLineToSubsection("HEAD", changeSectionLabel, changeLine)
	}

	return changes.String()
}

func deletableRef(pr *github.PullRequest, owner string) bool {
	return *pr.Head.Repo.Owner.Login == owner && *pr.Head.Ref != "master" && *pr.Head.Ref != "gh-pages"
}

func commitHistoryFile(context *ctx.Context, historySHA, owner, repo string, number int, newHistoryFileContents string) error {
	repositoryContentsOptions := &github.RepositoryContentFileOptions{
		Message: github.String(fmt.Sprintf("Update history to reflect merge of #%d [ci skip]", number)),
		Content: []byte(newHistoryFileContents),
		SHA:     github.String(historySHA),
		Committer: &github.CommitAuthor{
			Name:  github.String("jekyllbot"),
			Email: github.String("jekyllbot@jekyllrb.com"),
		},
	}
	updateResponse, _, err := context.GitHub.Repositories.UpdateFile(owner, repo, "History.markdown", repositoryContentsOptions)
	if err != nil {
		fmt.Printf("comments: error committing History.markdown: %v\n", err)
		return err
	}
	fmt.Printf("comments: updateResponse: %s\n", updateResponse)
	return nil
}
<commit_msg>Logic is backwards<commit_after>package issuecomment

import (
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"
	"sync"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/auth"
	"github.com/parkr/auto-reply/ctx"
	"github.com/parkr/changelog"
)

// changelogCategory is a changelog category, like "Site Enhancements" and such.
type changelogCategory struct {
	Prefix, Slug, Section string
	Labels                []string
}

var (
	mergeCommentRegexp = regexp.MustCompile("@[a-zA-Z-_]+: (merge|:shipit:|:ship:)( \\+([a-zA-Z-_ ]+))?")

	categories = []changelogCategory{
		changelogCategory{
			Prefix:  "major",
			Slug:    "major-enhancements",
			Section: "Major Enhancements",
			Labels:  []string{"feature"},
		},
		changelogCategory{
			Prefix:  "minor",
			Slug:    "minor-enhancements",
			Section: "Minor Enhancements",
			Labels:  []string{"enhancement"},
		},
		changelogCategory{
			Prefix:  "bug",
			Slug:    "bug-fixes",
			Section: "Bug Fixes",
			Labels:  []string{"bug", "fix"},
		},
		changelogCategory{
			Prefix:  "dev",
			Slug:    "development-fixes",
			Section: "Development Fixes",
			Labels:  []string{"internal", "fix"},
		},
		changelogCategory{
			Prefix:  "port",
			Slug:    "forward-ports",
			Section: "Forward Ports",
			Labels:  []string{"forward-port"},
		},
		changelogCategory{
			Prefix:  "site",
			Slug:    "site-enhancements",
			Section: "Site Enhancements",
			Labels:  []string{"documentation"},
		},
	}
)

func MergeAndLabel(context *ctx.Context, payload interface{}) error {
	event, ok := payload.(*github.IssueCommentEvent)
	if !ok {
		return context.NewError("MergeAndLabel: not an issue comment event")
	}

	// Is this a pull request?
	if event.Issue == nil || event.Issue.PullRequestLinks == nil {
		return context.NewError("MergeAndLabel: not a pull request")
	}

	var changeSectionLabel string
	isReq, labelFromComment := parseMergeRequestComment(*event.Comment.Body)

	// Is It a merge request comment?
	if !isReq {
		return context.NewError("MergeAndLabel: not a merge request comment")
	}

	if os.Getenv("AUTO_REPLY_DEBUG") == "true" {
		log.Println("MergeAndLabel: received event:", event)
	}

	var wg sync.WaitGroup

	owner := *event.Repo.Owner.Login
	repo := *event.Repo.Name
	number := *event.Issue.Number
	ref := fmt.Sprintf("%s/%s#%d", owner, repo, number)

	// Does the user have merge/label abilities?
	if !auth.CommenterHasPushAccess(context, *event) {
		log.Printf("%s isn't authenticated to merge anything on %s", *event.Comment.User.Login, *event.Repo.FullName)
		return errors.New("commenter isn't allowed to merge")
	}

	// Should it be labeled?
	if labelFromComment != "" {
		changeSectionLabel = sectionForLabel(labelFromComment)
	} else {
		changeSectionLabel = "none"
	}
	fmt.Printf("changeSectionLabel = '%s'\n", changeSectionLabel)

	// Merge
	commitMsg := fmt.Sprintf("Merge pull request %v", number)
	_, _, mergeErr := context.GitHub.PullRequests.Merge(owner, repo, number, commitMsg)
	if mergeErr != nil {
		return context.NewError("MergeAndLabel: error merging %s: %v", ref, mergeErr)
	}

	// Delete branch
	repoInfo, _, getRepoErr := context.GitHub.PullRequests.Get(owner, repo, number)
	if getRepoErr != nil {
		return context.NewError("MergeAndLabel: error getting PR info %s: %v", ref, getRepoErr)
	}

	// Delete branch
	if deletableRef(repoInfo, owner) {
		wg.Add(1)
		go func() {
			ref := fmt.Sprintf("heads/%s", *repoInfo.Head.Ref)
			_, deleteBranchErr := context.GitHub.Git.DeleteRef(owner, repo, ref)
			if deleteBranchErr != nil {
				fmt.Printf("MergeAndLabel: error deleting branch %v\n", mergeErr)
			}
			wg.Done()
		}()
	}

	wg.Add(1)
	go func() {
		err := addLabelsForSubsection(context, owner, repo, number, changeSectionLabel)
		if err != nil {
			fmt.Printf("MergeAndLabel: error applying labels: %v\n", err)
		}
		wg.Done()
	}()

	wg.Add(1)
	go func() {
		// Read History.markdown, add line to appropriate change section
		historyFileContents, historySHA := getHistoryContents(context, owner, repo)

		// Add merge reference to history
		newHistoryFileContents := addMergeReference(historyFileContents, changeSectionLabel, *repoInfo.Title, number)

		// Commit change to History.markdown
		commitErr := commitHistoryFile(context, historySHA, owner, repo, number, newHistoryFileContents)
		if commitErr != nil {
			fmt.Printf("comments: error committing updated history %v\n", mergeErr)
		}
		wg.Done()
	}()

	wg.Wait()

	return nil
}

func parseMergeRequestComment(commentBody string) (bool, string) {
	matches := mergeCommentRegexp.FindAllStringSubmatch(commentBody, -1)
	if matches == nil || matches[0] == nil {
		return false, ""
	}

	var label string
	if len(matches[0]) >= 4 {
		if labelFromComment := matches[0][3]; labelFromComment != "" {
			label = downcaseAndHyphenize(labelFromComment)
		}
	}

	return true, normalizeLabel(label)
}

func downcaseAndHyphenize(label string) string {
	return strings.Replace(strings.ToLower(label), " ", "-", -1)
}

func normalizeLabel(label string) string {
	for _, category := range categories {
		if strings.HasPrefix(label, category.Prefix) {
			return category.Slug
		}
	}

	return label
}

func sectionForLabel(slug string) string {
	for _, category := range categories {
		if slug == category.Slug {
			return category.Section
		}
	}

	return slug
}

func labelsForSubsection(changeSectionLabel string) []string {
	for _, category := range categories {
		if changeSectionLabel == category.Section {
			return category.Labels
		}
	}

	return []string{}
}

func selectSectionLabel(labels []github.Label) string {
	for _, label := range labels {
		if sectionForLabel(*label.Name) != *label.Name {
			return *label.Name
		}
	}
	return ""
}

func containsChangeLabel(commentBody string) bool {
	_, labelFromComment := parseMergeRequestComment(commentBody)
	return labelFromComment != ""
}

func addLabelsForSubsection(context *ctx.Context, owner, repo string, number int, changeSectionLabel string) error {
	labels := labelsForSubsection(changeSectionLabel)

	if len(labels) < 1 {
		return fmt.Errorf("no labels for changeSectionLabel='%s'", changeSectionLabel)
	}

	_, _, err := context.GitHub.Issues.AddLabelsToIssue(owner, repo, number, labels)
	return err
}

func getHistoryContents(context *ctx.Context, owner, repo string) (content, sha string) {
	contents, _, _, err := context.GitHub.Repositories.GetContents(
		owner,
		repo,
		"History.markdown",
		&github.RepositoryContentGetOptions{Ref: "heads/master"},
	)
	if err != nil {
		fmt.Printf("comments: error getting History.markdown %v\n", err)
		return "", ""
	}
	return base64Decode(*contents.Content), *contents.SHA
}

func base64Decode(encoded string) string {
	decoded, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		fmt.Printf("comments: error decoding string: %v\n", err)
		return ""
	}
	return string(decoded)
}

func addMergeReference(historyFileContents, changeSectionLabel, prTitle string, number int) string {
	changes, err := changelog.NewChangelogFromReader(strings.NewReader(historyFileContents))
	if historyFileContents == "" {
		err = nil
		changes = changelog.NewChangelog()
	}
	if err != nil {
		fmt.Printf("comments: error %v\n", err)
		return historyFileContents
	}

	changeLine := &changelog.ChangeLine{
		Summary:   prTitle,
		Reference: fmt.Sprintf("#%d", number),
	}

	// Put either directly in the version history or in a subsection.
	if changeSectionLabel == "none" {
		changes.AddLineToVersion("HEAD", changeLine)
	} else {
		changes.AddLineToSubsection("HEAD", changeSectionLabel, changeLine)
	}

	return changes.String()
}

func deletableRef(pr *github.PullRequest, owner string) bool {
	return *pr.Head.Repo.Owner.Login == owner && *pr.Head.Ref != "master" && *pr.Head.Ref != "gh-pages"
}

func commitHistoryFile(context *ctx.Context, historySHA, owner, repo string, number int, newHistoryFileContents string) error {
	repositoryContentsOptions := &github.RepositoryContentFileOptions{
		Message: github.String(fmt.Sprintf("Update history to reflect merge of #%d [ci skip]", number)),
		Content: []byte(newHistoryFileContents),
		SHA:     github.String(historySHA),
		Committer: &github.CommitAuthor{
			Name:  github.String("jekyllbot"),
			Email: github.String("jekyllbot@jekyllrb.com"),
		},
	}
	updateResponse, _, err := context.GitHub.Repositories.UpdateFile(owner, repo, "History.markdown", repositoryContentsOptions)
	if err != nil {
		fmt.Printf("comments: error committing History.markdown: %v\n", err)
		return err
	}
	fmt.Printf("comments: updateResponse: %s\n", updateResponse)
	return nil
}
<|endoftext|>


<commit_before>package driver

import ()

// ID プロバイダ選択時に列挙する情報。
type IdProvider struct {
	Uuid     string `json:"uuid"       bson:"uuid"`
	Name     string `json:"name"       bson:"name"`
	LoginUri string `json:"login_uri"  bson:"login_uri"`
}

// ID プロバイダの列挙。
type IdProviderLister interface {
	IdProviders() ([]*IdProvider, error)
}

// ID プロバイダの列挙。キャッシュ用。
type DatedIdProviderLister interface {
	StampedIdProviders(caStmp *Stamp) ([]*IdProvider, *Stamp, error)
}
<commit_msg>URL パラメータの引き継ぎ処理の改善<commit_after>package driver

import ()

// ID プロバイダ選択時に列挙する情報。
type IdProvider struct {
	Uuid     string `json:"uuid"       bson:"uuid"`
	Name     string `json:"name"       bson:"name"`
	LoginUri string `json:"login_uri"  bson:"login_uri"`
}

func (idp *IdProvider) String() string {
	return idp.Uuid + "," + idp.Name + "," + idp.LoginUri
}

// ID プロバイダの列挙。
type IdProviderLister interface {
	IdProviders() ([]*IdProvider, error)
}

// ID プロバイダの列挙。キャッシュ用。
type DatedIdProviderLister interface {
	StampedIdProviders(caStmp *Stamp) ([]*IdProvider, *Stamp, error)
}
<|endoftext|>


<commit_before>package misc

import (
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"
)

func TestGetProgramMemUsageKB(t *testing.T) {
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetProgramMemoryUsageKB()
		return
	}
	if usage := GetProgramMemoryUsageKB(); usage < 1000 {
		t.Fatal(usage)
	}
}

func TestGetSystemMemoryUsageKB(t *testing.T) {
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetSystemMemoryUsageKB()
		return
	}
	used, total := GetSystemMemoryUsageKB()
	if used < 1000 || total < used {
		t.Fatal(used, total)
	}
}

func TestGetSystemLoad(t *testing.T) {
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetSystemMemoryUsageKB()
		return
	}
	load := GetSystemLoad()
	if len(load) < 6 {
		t.Fatal(load)
	}
}

func TestGetSystemUptimeSec(t *testing.T) {
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetSystemUptimeSec()
		return
	}
	uptime := GetSystemUptimeSec()
	if uptime < 10 {
		t.Fatal(uptime)
	}
}

func TestGetRootDiskUsageKB(t *testing.T) {
	if HostIsWindows() || HostIsCircleCI() {
		// Just make sure the function does not crash
		GetRootDiskUsageKB()
		return
	}
	used, free, total := GetRootDiskUsageKB()
	if used == 0 || free == 0 || total == 0 || used+free != total {
		t.Fatal(used/1024, free/1024, total/1024)
	}
}

func TestGetSysctl(t *testing.T) {
	key := "kernel.pid_max"
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetSysctlInt(key)
		GetSysctlStr(key)
		return
	}
	if val, err := GetSysctlStr(key); err != nil || val == "" {
		t.Fatal(val, err)
	}
	if val, err := GetSysctlInt(key); err != nil || val < 1 {
		t.Fatal(val, err)
	}
	if old, err := IncreaseSysctlInt(key, 65535); old == 0 ||
		(err != nil && !strings.Contains(err.Error(), "permission") && !strings.Contains(err.Error(), "read-only")) {
		t.Fatal(err)
	}
}

func TestInvokeProgram(t *testing.T) {
	if HostIsWindows() {
		out, err := InvokeProgram([]string{"A=laitos123"}, 3, "hostname")
		if err != nil || len(out) < 1 {
			t.Fatal(err, out)
		}

		begin := time.Now()
		out, err = InvokeProgram(nil, 3, "cmd.exe", "/c", "waitfor dummydummy /t 60")
		if err == nil {
			t.Fatal("did not timeout")
		}
		duration := time.Now().Unix() - begin.Unix()
		if duration > 4 {
			t.Fatal("did not kill before timeout")
		}
	} else {
		out, err := InvokeProgram([]string{"A=laitos123"}, 10, "printenv", "A")
		if err != nil || out != "laitos123\n" {
			t.Fatal(err, out)
		}

		begin := time.Now()
		out, err = InvokeProgram(nil, 1, "sleep", "5")
		if err == nil {
			t.Fatal("did not timeout")
		}
		duration := time.Now().Unix() - begin.Unix()
		if duration > 2 {
			t.Fatal("did not kill before timeout")
		}
	}
}

func TestInvokeShell(t *testing.T) {
	if HostIsWindows() {
		out, err := InvokeShell(3, "C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe", "echo $env:windir")
		if err != nil || !strings.Contains(strings.ToLower(out), "windows") {
			t.Fatal(err, out)
		}
	} else {
		out, err := InvokeShell(1, "/bin/sh", "echo $PATH")
		if err != nil || out != CommonPATH+"\n" {
			t.Fatal(err, out)
		}
	}
}

func TestPrepareUtilities(t *testing.T) {
	PrepareUtilities(Logger{})
	for _, utilName := range []string{"busybox", "toybox", "phantomjs"} {
		if _, err := os.Stat(filepath.Join(UtilityDir, utilName)); err != nil {
			t.Fatal("cannot find program "+utilName, err)
		}
	}
}

func TestGetLocalUserNames(t *testing.T) {
	names := GetLocalUserNames()
	if len(names) < 2 || (!names["root"] && !names["Administrator"]) {
		t.Fatal(names)
	}
}

func TestBlockUserLogin(t *testing.T) {
	// just make sure it does not panic
	t.Log(BlockUserLogin("nobody"))
}

func TestDisableStopDaemon(t *testing.T) {
	// just make sure it does not panic
	t.Log(DisableStopDaemon("this-service-does-not-exist"))
}

func TestEnableStartDaemon(t *testing.T) {
	// just make sure it does not panic
	t.Log(EnableStartDaemon("this-service-does-not-exist"))
}

func TestDisableInterferingResolvd(t *testing.T) {
	// just make sure it does not panic
	t.Log(DisableInterferingResolved())
}

func TestSwapOff(t *testing.T) {
	// just make sure it does not panic
	SwapOff()
}

func TestSetTermEcho(t *testing.T) {
	// just make sure it does not panic
	SetTermEcho(false)
	SetTermEcho(true)
}

func TestLockMemory(t *testing.T) {
	// just make sure it does not panic
	LockMemory()
}

func TestSetTimeZone(t *testing.T) {
	// just make sure it does not panic
	t.Log(SetTimeZone("UTC"))
}
<commit_msg>skip PrepareUtilities test on Windows<commit_after>package misc

import (
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"
)

func TestGetProgramMemUsageKB(t *testing.T) {
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetProgramMemoryUsageKB()
		return
	}
	if usage := GetProgramMemoryUsageKB(); usage < 1000 {
		t.Fatal(usage)
	}
}

func TestGetSystemMemoryUsageKB(t *testing.T) {
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetSystemMemoryUsageKB()
		return
	}
	used, total := GetSystemMemoryUsageKB()
	if used < 1000 || total < used {
		t.Fatal(used, total)
	}
}

func TestGetSystemLoad(t *testing.T) {
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetSystemMemoryUsageKB()
		return
	}
	load := GetSystemLoad()
	if len(load) < 6 {
		t.Fatal(load)
	}
}

func TestGetSystemUptimeSec(t *testing.T) {
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetSystemUptimeSec()
		return
	}
	uptime := GetSystemUptimeSec()
	if uptime < 10 {
		t.Fatal(uptime)
	}
}

func TestGetRootDiskUsageKB(t *testing.T) {
	if HostIsWindows() || HostIsCircleCI() {
		// Just make sure the function does not crash
		GetRootDiskUsageKB()
		return
	}
	used, free, total := GetRootDiskUsageKB()
	if used == 0 || free == 0 || total == 0 || used+free != total {
		t.Fatal(used/1024, free/1024, total/1024)
	}
}

func TestGetSysctl(t *testing.T) {
	key := "kernel.pid_max"
	if runtime.GOOS != "linux" {
		// Just make sure the function does not crash
		GetSysctlInt(key)
		GetSysctlStr(key)
		return
	}
	if val, err := GetSysctlStr(key); err != nil || val == "" {
		t.Fatal(val, err)
	}
	if val, err := GetSysctlInt(key); err != nil || val < 1 {
		t.Fatal(val, err)
	}
	if old, err := IncreaseSysctlInt(key, 65535); old == 0 ||
		(err != nil && !strings.Contains(err.Error(), "permission") && !strings.Contains(err.Error(), "read-only")) {
		t.Fatal(err)
	}
}

func TestInvokeProgram(t *testing.T) {
	if HostIsWindows() {
		out, err := InvokeProgram([]string{"A=laitos123"}, 3, "hostname")
		if err != nil || len(out) < 1 {
			t.Fatal(err, out)
		}

		begin := time.Now()
		out, err = InvokeProgram(nil, 3, "cmd.exe", "/c", "waitfor dummydummy /t 60")
		if err == nil {
			t.Fatal("did not timeout")
		}
		duration := time.Now().Unix() - begin.Unix()
		if duration > 4 {
			t.Fatal("did not kill before timeout")
		}
	} else {
		out, err := InvokeProgram([]string{"A=laitos123"}, 10, "printenv", "A")
		if err != nil || out != "laitos123\n" {
			t.Fatal(err, out)
		}

		begin := time.Now()
		out, err = InvokeProgram(nil, 1, "sleep", "5")
		if err == nil {
			t.Fatal("did not timeout")
		}
		duration := time.Now().Unix() - begin.Unix()
		if duration > 2 {
			t.Fatal("did not kill before timeout")
		}
	}
}

func TestInvokeShell(t *testing.T) {
	if HostIsWindows() {
		out, err := InvokeShell(3, "C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe", "echo $env:windir")
		if err != nil || !strings.Contains(strings.ToLower(out), "windows") {
			t.Fatal(err, out)
		}
	} else {
		out, err := InvokeShell(1, "/bin/sh", "echo $PATH")
		if err != nil || out != CommonPATH+"\n" {
			t.Fatal(err, out)
		}
	}
}

func TestPrepareUtilities(t *testing.T) {
	PrepareUtilities(Logger{})
	if HostIsWindows() {
		// Just make sure it does not panic
		return
	}
	for _, utilName := range []string{"busybox", "toybox", "phantomjs"} {
		if _, err := os.Stat(filepath.Join(UtilityDir, utilName)); err != nil {
			t.Fatal("cannot find program "+utilName, err)
		}
	}
}

func TestGetLocalUserNames(t *testing.T) {
	names := GetLocalUserNames()
	if len(names) < 2 || (!names["root"] && !names["Administrator"]) {
		t.Fatal(names)
	}
}

func TestBlockUserLogin(t *testing.T) {
	// just make sure it does not panic
	t.Log(BlockUserLogin("nobody"))
}

func TestDisableStopDaemon(t *testing.T) {
	// just make sure it does not panic
	t.Log(DisableStopDaemon("this-service-does-not-exist"))
}

func TestEnableStartDaemon(t *testing.T) {
	// just make sure it does not panic
	t.Log(EnableStartDaemon("this-service-does-not-exist"))
}

func TestDisableInterferingResolvd(t *testing.T) {
	// just make sure it does not panic
	t.Log(DisableInterferingResolved())
}

func TestSwapOff(t *testing.T) {
	// just make sure it does not panic
	SwapOff()
}

func TestSetTermEcho(t *testing.T) {
	// just make sure it does not panic
	SetTermEcho(false)
	SetTermEcho(true)
}

func TestLockMemory(t *testing.T) {
	// just make sure it does not panic
	LockMemory()
}

func TestSetTimeZone(t *testing.T) {
	// just make sure it does not panic
	t.Log(SetTimeZone("UTC"))
}
<|endoftext|>


<commit_before>// Package mmhook implements an adapter that uses Mattermost Webhooks.
package mmhook

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/gorilla/schema"
	"github.com/yukithm/mmbot/adapter"
	"github.com/yukithm/mmbot/message"
)

// Client is a client for Mattermost.
type Client struct {
	config *adapter.Config
	logger *log.Logger
	http   *http.Client
	tokens map[string]int
	in     chan message.InMessage
	errCh  chan error
}

// NewClient returns new mattermost webhook client.
func NewClient(config *adapter.Config, logger *log.Logger) *Client {
	if logger == nil {
		logger = log.New(ioutil.Discard, "", 0)
	}
	c := &Client{
		config: config,
		logger: logger,
	}

	tr := &http.Transport{
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	if config.InsecureSkipVerify {
		tr.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: config.InsecureSkipVerify,
		}
	}
	c.http = &http.Client{Transport: tr}

	// build token lookup table
	c.tokens = make(map[string]int, len(config.Tokens))
	for i, token := range config.Tokens {
		token = strings.TrimSpace(token)
		if token != "" {
			c.tokens[token] = i
		}
	}

	return c
}

// Start starts the communication with Mattermost.
func (c *Client) Start() (chan message.InMessage, chan error) {
	c.in = make(chan message.InMessage, 1)
	c.errCh = make(chan error, 1)
	return c.in, c.errCh
}

// Stop terminates the communication.
func (c *Client) Stop() {
	close(c.in)
	close(c.errCh)
}

// Send sends a message to Mattermost.
func (c *Client) Send(msg *message.OutMessage) error {
	om := translateOutMessage(msg)
	if c.config.OverrideUserName != "" && om.UserName == "" {
		om.UserName = c.config.OverrideUserName
	}
	if c.config.IconURL != "" && om.IconURL == "" {
		om.IconURL = c.config.IconURL
	}

	buf, err := json.Marshal(om)
	if err != nil {
		return err
	}

	res, err := c.http.Post(c.config.OutgoingURL, "application/json", bytes.NewReader(buf))
	if err != nil {
		return err
	}
	defer res.Body.Close()

	io.Copy(ioutil.Discard, res.Body)
	if res.StatusCode != 200 {
		return fmt.Errorf("Failed to send a message (%d %s)",
			res.StatusCode, res.Status)
	}

	return nil
}

// IncomingWebHook returns webhook. It will be disabled if nil.
func (c *Client) IncomingWebHook() *adapter.IncomingWebHook {
	return &adapter.IncomingWebHook{
		Path:    c.config.IncomingPath,
		Handler: c.ServeHTTP,
	}
}

// ServeHTTP implements http.Handler interface.
// ServeHTTP receives a message from Mattermost outgoing webhook.
func (c *Client) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		c.logger.Printf("Invalid %q request from %q", r.Method, r.RemoteAddr)
		http.Error(w, "405 Method Not Allowed", http.StatusMethodNotAllowed)
		return
	}

	msg := InMessage{}
	if err := decodeForm(&msg, r); err != nil {
		c.logger.Printf("Invalid form data: %v", err)
		http.Error(w, "400 Bad Request", http.StatusBadRequest)
		return
	}

	if len(c.tokens) > 0 {
		if msg.Token == "" {
			c.logger.Printf("No token request from %q", r.RemoteAddr)
			http.Error(w, "400 Bad Request", http.StatusBadRequest)
			return
		} else if !c.validToken(msg.Token) {
			c.logger.Printf("Invalid token %q request from %q", msg.Token, r.RemoteAddr)
			http.Error(w, "400 Bad Request", http.StatusBadRequest)
			return
		}
	}

	im := translateInMessage(&msg)
	c.in <- *im
}

func (c *Client) validToken(token string) bool {
	_, ok := c.tokens[token]
	return ok
}

func decodeForm(msg *InMessage, r *http.Request) error {
	if err := r.ParseForm(); err != nil {
		return err
	}

	decoder := schema.NewDecoder()
	if err := decoder.Decode(msg, r.PostForm); err != nil {
		return err
	}

	return nil
}
<commit_msg>Return more detailed error from mmhook.Client.Send()<commit_after>// Package mmhook implements an adapter that uses Mattermost Webhooks.
package mmhook

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/schema"
	"github.com/yukithm/mmbot/adapter"
	"github.com/yukithm/mmbot/message"
)

// SendError represents an error of sending message.
type SendError struct {
	Err                error
	RatelimitLimit     int
	RatelimitRemaining int
	RatelimitReset     int
	RequestID          string
	VersionID          string
	Date               string
	ContentLength      int
	ContentType        string
	Body               string
}

func (e SendError) Error() string {
	return e.Err.Error()
}

// Client is a client for Mattermost.
type Client struct {
	config *adapter.Config
	logger *log.Logger
	http   *http.Client
	tokens map[string]int
	in     chan message.InMessage
	errCh  chan error
}

// NewClient returns new mattermost webhook client.
func NewClient(config *adapter.Config, logger *log.Logger) *Client {
	if logger == nil {
		logger = log.New(ioutil.Discard, "", 0)
	}
	c := &Client{
		config: config,
		logger: logger,
	}

	tr := &http.Transport{
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	if config.InsecureSkipVerify {
		tr.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: config.InsecureSkipVerify,
		}
	}
	c.http = &http.Client{Transport: tr}

	// build token lookup table
	c.tokens = make(map[string]int, len(config.Tokens))
	for i, token := range config.Tokens {
		token = strings.TrimSpace(token)
		if token != "" {
			c.tokens[token] = i
		}
	}

	return c
}

// Start starts the communication with Mattermost.
func (c *Client) Start() (chan message.InMessage, chan error) {
	c.in = make(chan message.InMessage, 1)
	c.errCh = make(chan error, 1)
	return c.in, c.errCh
}

// Stop terminates the communication.
func (c *Client) Stop() {
	close(c.in)
	close(c.errCh)
}

// Send sends a message to Mattermost.
func (c *Client) Send(msg *message.OutMessage) error {
	om := translateOutMessage(msg)
	if c.config.OverrideUserName != "" && om.UserName == "" {
		om.UserName = c.config.OverrideUserName
	}
	if c.config.IconURL != "" && om.IconURL == "" {
		om.IconURL = c.config.IconURL
	}

	buf, err := json.Marshal(om)
	if err != nil {
		return err
	}

	res, err := c.http.Post(c.config.OutgoingURL, "application/json", bytes.NewReader(buf))
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode == 200 {
		io.Copy(ioutil.Discard, res.Body)
	} else {
		return newSendError(res)
	}

	return nil
}

// IncomingWebHook returns webhook. It will be disabled if nil.
func (c *Client) IncomingWebHook() *adapter.IncomingWebHook {
	return &adapter.IncomingWebHook{
		Path:    c.config.IncomingPath,
		Handler: c.ServeHTTP,
	}
}

// ServeHTTP implements http.Handler interface.
// ServeHTTP receives a message from Mattermost outgoing webhook.
func (c *Client) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		c.logger.Printf("Invalid %q request from %q", r.Method, r.RemoteAddr)
		http.Error(w, "405 Method Not Allowed", http.StatusMethodNotAllowed)
		return
	}

	msg := InMessage{}
	if err := decodeForm(&msg, r); err != nil {
		c.logger.Printf("Invalid form data: %v", err)
		http.Error(w, "400 Bad Request", http.StatusBadRequest)
		return
	}

	if len(c.tokens) > 0 {
		if msg.Token == "" {
			c.logger.Printf("No token request from %q", r.RemoteAddr)
			http.Error(w, "400 Bad Request", http.StatusBadRequest)
			return
		} else if !c.validToken(msg.Token) {
			c.logger.Printf("Invalid token %q request from %q", msg.Token, r.RemoteAddr)
			http.Error(w, "400 Bad Request", http.StatusBadRequest)
			return
		}
	}

	im := translateInMessage(&msg)
	c.in <- *im
}

func (c *Client) validToken(token string) bool {
	_, ok := c.tokens[token]
	return ok
}

func decodeForm(msg *InMessage, r *http.Request) error {
	if err := r.ParseForm(); err != nil {
		return err
	}

	decoder := schema.NewDecoder()
	if err := decoder.Decode(msg, r.PostForm); err != nil {
		return err
	}

	return nil
}

func newSendError(res *http.Response) SendError {
	var body bytes.Buffer
	io.Copy(&body, res.Body)
	return SendError{
		Err:                fmt.Errorf("Failed to send a message (%s)", res.Status),
		RatelimitLimit:     getHeaderInt(res.Header, "X-Ratelimit-Limit"),
		RatelimitRemaining: getHeaderInt(res.Header, "X-Ratelimit-Remaining"),
		RatelimitReset:     getHeaderInt(res.Header, "X-Ratelimit-Reset"),
		RequestID:          res.Header.Get("X-Request-Id"),
		VersionID:          res.Header.Get("X-Version-Id"),
		Date:               res.Header.Get("Date"),
		ContentLength:      getHeaderInt(res.Header, "Content-Length"),
		ContentType:        res.Header.Get("Content-Type"),
		Body:               body.String(),
	}
}

func getHeaderInt(header http.Header, key string) int {
	value := header.Get(key)
	if value != "" {
		n, err := strconv.Atoi(value)
		if err != nil {
			return -1
		}
		return n
	}
	return -1
}
<|endoftext|>


<commit_before>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"io"
	"path/filepath"
	"sort"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/super_block"
)

func init() {
	Commands = append(Commands, &commandVolumeFixReplication{})
}

type commandVolumeFixReplication struct {
	collectionPattern *string
}

func (c *commandVolumeFixReplication) Name() string {
	return "volume.fix.replication"
}

func (c *commandVolumeFixReplication) Help() string {
	return `add replicas to volumes that are missing replicas

	This command finds all over-replicated volumes. If found, it will purge the oldest copies and stop.

	This command also finds all under-replicated volumes, and finds volume servers with free slots.
	If the free slots satisfy the replication requirement, the volume content is copied over and mounted.

	volume.fix.replication -n                             # do not take action
	volume.fix.replication                                # actually deleting or copying the volume files and mount the volume
	volume.fix.replication -collectionPattern=important*  # fix any collections with prefix "important"

	Note:
		* each time this will only add back one replica for one volume id. If there are multiple replicas
		  are missing, e.g. multiple volume servers are new, you may need to run this multiple times.
		* do not run this too quickly within seconds, since the new volume replica may take a few seconds
		  to register itself to the master.

`
}

func (c *commandVolumeFixReplication) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	volFixReplicationCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	c.collectionPattern = volFixReplicationCommand.String("collectionPattern", "", "match with wildcard characters '*' and '?'")
	skipChange := volFixReplicationCommand.Bool("n", false, "skip the changes")
	if err = volFixReplicationCommand.Parse(args); err != nil {
		return nil
	}

	takeAction := !*skipChange

	var resp *master_pb.VolumeListResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return err
	}

	// find all volumes that needs replication
	// collect all data nodes
	volumeReplicas, allLocations := collectVolumeReplicaLocations(resp)

	if len(allLocations) == 0 {
		return fmt.Errorf("no data nodes at all")
	}

	// find all under replicated volumes
	var underReplicatedVolumeIds, overReplicatedVolumeIds []uint32
	for vid, replicas := range volumeReplicas {
		replica := replicas[0]
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(replica.info.ReplicaPlacement))
		if replicaPlacement.GetCopyCount() > len(replicas) {
			underReplicatedVolumeIds = append(underReplicatedVolumeIds, vid)
		} else if replicaPlacement.GetCopyCount() < len(replicas) {
			overReplicatedVolumeIds = append(overReplicatedVolumeIds, vid)
			fmt.Fprintf(writer, "volume %d replication %s, but over replicated %+d\n", replica.info.Id, replicaPlacement, len(replicas))
		}
	}

	if len(overReplicatedVolumeIds) > 0 {
		return c.fixOverReplicatedVolumes(commandEnv, writer, takeAction, overReplicatedVolumeIds, volumeReplicas, allLocations)
	}

	if len(underReplicatedVolumeIds) == 0 {
		return nil
	}

	// find the most under populated data nodes
	keepDataNodesSorted(allLocations)

	return c.fixUnderReplicatedVolumes(commandEnv, writer, takeAction, underReplicatedVolumeIds, volumeReplicas, allLocations)

}

func collectVolumeReplicaLocations(resp *master_pb.VolumeListResponse) (map[uint32][]*VolumeReplica, []location) {
	volumeReplicas := make(map[uint32][]*VolumeReplica)
	var allLocations []location
	eachDataNode(resp.TopologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		loc := newLocation(dc, string(rack), dn)
		for _, v := range dn.VolumeInfos {
			volumeReplicas[v.Id] = append(volumeReplicas[v.Id], &VolumeReplica{
				location: &loc,
				info:     v,
			})
		}
		allLocations = append(allLocations, loc)
	})
	return volumeReplicas, allLocations
}

func (c *commandVolumeFixReplication) fixOverReplicatedVolumes(commandEnv *CommandEnv, writer io.Writer, takeAction bool, overReplicatedVolumeIds []uint32, volumeReplicas map[uint32][]*VolumeReplica, allLocations []location) error {
	for _, vid := range overReplicatedVolumeIds {
		replicas := volumeReplicas[vid]
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(replicas[0].info.ReplicaPlacement))

		replica := pickOneReplicaToDelete(replicas, replicaPlacement)

		// check collection name pattern
		if *c.collectionPattern != "" {
			matched, err := filepath.Match(*c.collectionPattern, replica.info.Collection)
			if err != nil {
				return fmt.Errorf("match pattern %s with collection %s: %v", *c.collectionPattern, replica.info.Collection, err)
			}
			if !matched {
				break
			}
		}

		fmt.Fprintf(writer, "deleting volume %d from %s ...\n", replica.info.Id, replica.location.dataNode.Id)

		if !takeAction {
			break
		}

		if err := deleteVolume(commandEnv.option.GrpcDialOption, needle.VolumeId(replica.info.Id), replica.location.dataNode.Id); err != nil {
			return fmt.Errorf("deleting volume %d from %s : %v", replica.info.Id, replica.location.dataNode.Id, err)
		}

	}
	return nil
}

func (c *commandVolumeFixReplication) fixUnderReplicatedVolumes(commandEnv *CommandEnv, writer io.Writer, takeAction bool, underReplicatedVolumeIds []uint32, volumeReplicas map[uint32][]*VolumeReplica, allLocations []location) error {
	for _, vid := range underReplicatedVolumeIds {
		replicas := volumeReplicas[vid]
		replica := pickOneReplicaToCopyFrom(replicas)
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(replica.info.ReplicaPlacement))
		foundNewLocation := false
		for _, dst := range allLocations {
			// check whether data nodes satisfy the constraints
			if dst.dataNode.FreeVolumeCount > 0 && satisfyReplicaPlacement(replicaPlacement, replicas, dst) {
				// check collection name pattern
				if *c.collectionPattern != "" {
					matched, err := filepath.Match(*c.collectionPattern, replica.info.Collection)
					if err != nil {
						return fmt.Errorf("match pattern %s with collection %s: %v", *c.collectionPattern, replica.info.Collection, err)
					}
					if !matched {
						break
					}
				}

				// ask the volume server to replicate the volume
				foundNewLocation = true
				fmt.Fprintf(writer, "replicating volume %d %s from %s to dataNode %s ...\n", replica.info.Id, replicaPlacement, replica.location.dataNode.Id, dst.dataNode.Id)

				if !takeAction {
					break
				}

				err := operation.WithVolumeServerClient(dst.dataNode.Id, commandEnv.option.GrpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
					_, replicateErr := volumeServerClient.VolumeCopy(context.Background(), &volume_server_pb.VolumeCopyRequest{
						VolumeId:       replica.info.Id,
						SourceDataNode: replica.location.dataNode.Id,
					})
					if replicateErr != nil {
						return fmt.Errorf("copying from %s => %s : %v", replica.location.dataNode.Id, dst.dataNode.Id, replicateErr)
					}
					return nil
				})

				if err != nil {
					return err
				}

				// adjust free volume count
				dst.dataNode.FreeVolumeCount--
				keepDataNodesSorted(allLocations)
				break
			}
		}
		if !foundNewLocation {
			fmt.Fprintf(writer, "failed to place volume %d replica as %s, existing:%+v\n", replica.info.Id, replicaPlacement, len(replicas))
		}

	}
	return nil
}

func keepDataNodesSorted(dataNodes []location) {
	sort.Slice(dataNodes, func(i, j int) bool {
		return dataNodes[i].dataNode.FreeVolumeCount > dataNodes[j].dataNode.FreeVolumeCount
	})
}

/*
  if on an existing data node {
    return false
  }
  if different from existing dcs {
    if lack on different dcs {
      return true
    }else{
      return false
    }
  }
  if not on primary dc {
    return false
  }
  if different from existing racks {
    if lack on different racks {
      return true
    }else{
      return false
    }
  }
  if not on primary rack {
    return false
  }
  if lacks on same rack {
    return true
  } else {
    return false
  }
*/
func satisfyReplicaPlacement(replicaPlacement *super_block.ReplicaPlacement, replicas []*VolumeReplica, possibleLocation location) bool {

	existingDataCenters, _, existingDataNodes := countReplicas(replicas)

	if _, found := existingDataNodes[possibleLocation.String()]; found {
		// avoid duplicated volume on the same data node
		return false
	}

	primaryDataCenters, _ := findTopKeys(existingDataCenters)

	// ensure data center count is within limit
	if _, found := existingDataCenters[possibleLocation.DataCenter()]; !found {
		// different from existing dcs
		if len(existingDataCenters) < replicaPlacement.DiffDataCenterCount+1 {
			// lack on different dcs
			return true
		} else {
			// adding this would go over the different dcs limit
			return false
		}
	}
	// now this is same as one of the existing data center
	if !isAmong(possibleLocation.DataCenter(), primaryDataCenters) {
		// not on one of the primary dcs
		return false
	}

	// now this is one of the primary dcs
	primaryDcRacks := make(map[string]int)
	for _, replica := range replicas {
		if replica.location.DataCenter() != possibleLocation.DataCenter() {
			continue
		}
		primaryDcRacks[replica.location.Rack()] += 1
	}
	primaryRacks, _ := findTopKeys(primaryDcRacks)
	sameRackCount := primaryDcRacks[possibleLocation.Rack()]

	// ensure rack count is within limit
	if _, found := primaryDcRacks[possibleLocation.Rack()]; !found {
		// different from existing racks
		if len(primaryDcRacks) < replicaPlacement.DiffRackCount+1 {
			// lack on different racks
			return true
		} else {
			// adding this would go over the different racks limit
			return false
		}
	}
	// now this is same as one of the existing racks
	if !isAmong(possibleLocation.Rack(), primaryRacks) {
		// not on the primary rack
		return false
	}

	// now this is on the primary rack

	// different from existing data nodes
	if sameRackCount < replicaPlacement.SameRackCount+1 {
		// lack on same rack
		return true
	} else {
		// adding this would go over the same data node limit
		return false
	}

}

func findTopKeys(m map[string]int) (topKeys []string, max int) {
	for k, c := range m {
		if max < c {
			topKeys = topKeys[:0]
			topKeys = append(topKeys, k)
			max = c
		} else if max == c {
			topKeys = append(topKeys, k)
		}
	}
	return
}

func isAmong(key string, keys []string) bool {
	for _, k := range keys {
		if k == key {
			return true
		}
	}
	return false
}

type VolumeReplica struct {
	location *location
	info     *master_pb.VolumeInformationMessage
}

type location struct {
	dc       string
	rack     string
	dataNode *master_pb.DataNodeInfo
}

func newLocation(dc, rack string, dataNode *master_pb.DataNodeInfo) location {
	return location{
		dc:       dc,
		rack:     rack,
		dataNode: dataNode,
	}
}

func (l location) String() string {
	return fmt.Sprintf("%s %s %s", l.dc, l.rack, l.dataNode.Id)
}

func (l location) Rack() string {
	return fmt.Sprintf("%s %s", l.dc, l.rack)
}

func (l location) DataCenter() string {
	return l.dc
}

func pickOneReplicaToCopyFrom(replicas []*VolumeReplica) *VolumeReplica {
	mostRecent := replicas[0]
	for _, replica := range replicas {
		if replica.info.ModifiedAtSecond > mostRecent.info.ModifiedAtSecond {
			mostRecent = replica
		}
	}
	return mostRecent
}

func countReplicas(replicas []*VolumeReplica) (diffDc, diffRack, diffNode map[string]int) {
	diffDc = make(map[string]int)
	diffRack = make(map[string]int)
	diffNode = make(map[string]int)
	for _, replica := range replicas {
		diffDc[replica.location.DataCenter()] += 1
		diffRack[replica.location.Rack()] += 1
		diffNode[replica.location.String()] += 1
	}
	return
}

func pickOneReplicaToDelete(replicas []*VolumeReplica, replicaPlacement *super_block.ReplicaPlacement) *VolumeReplica {

	sort.Slice(replicas, func(i, j int) bool {
		a, b := replicas[i], replicas[j]
		if a.info.CompactRevision != b.info.CompactRevision {
			return a.info.CompactRevision < b.info.CompactRevision
		}
		if a.info.ModifiedAtSecond != b.info.ModifiedAtSecond {
			return a.info.ModifiedAtSecond < b.info.ModifiedAtSecond
		}
		if a.info.Size != b.info.Size {
			return a.info.Size < b.info.Size
		}
		return false
	})

	return replicas[0]

}
<commit_msg>avoid error message<commit_after>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"io"
	"path/filepath"
	"sort"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/super_block"
)

func init() {
	Commands = append(Commands, &commandVolumeFixReplication{})
}

type commandVolumeFixReplication struct {
	collectionPattern *string
}

func (c *commandVolumeFixReplication) Name() string {
	return "volume.fix.replication"
}

func (c *commandVolumeFixReplication) Help() string {
	return `add replicas to volumes that are missing replicas

	This command finds all over-replicated volumes. If found, it will purge the oldest copies and stop.

	This command also finds all under-replicated volumes, and finds volume servers with free slots.
	If the free slots satisfy the replication requirement, the volume content is copied over and mounted.

	volume.fix.replication -n                             # do not take action
	volume.fix.replication                                # actually deleting or copying the volume files and mount the volume
	volume.fix.replication -collectionPattern=important*  # fix any collections with prefix "important"

	Note:
		* each time this will only add back one replica for one volume id. If there are multiple replicas
		  are missing, e.g. multiple volume servers are new, you may need to run this multiple times.
		* do not run this too quickly within seconds, since the new volume replica may take a few seconds
		  to register itself to the master.

`
}

func (c *commandVolumeFixReplication) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	volFixReplicationCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	c.collectionPattern = volFixReplicationCommand.String("collectionPattern", "", "match with wildcard characters '*' and '?'")
	skipChange := volFixReplicationCommand.Bool("n", false, "skip the changes")
	if err = volFixReplicationCommand.Parse(args); err != nil {
		return nil
	}

	takeAction := !*skipChange

	var resp *master_pb.VolumeListResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return err
	}

	// find all volumes that needs replication
	// collect all data nodes
	volumeReplicas, allLocations := collectVolumeReplicaLocations(resp)

	if len(allLocations) == 0 {
		return fmt.Errorf("no data nodes at all")
	}

	// find all under replicated volumes
	var underReplicatedVolumeIds, overReplicatedVolumeIds []uint32
	for vid, replicas := range volumeReplicas {
		replica := replicas[0]
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(replica.info.ReplicaPlacement))
		if replicaPlacement.GetCopyCount() > len(replicas) {
			underReplicatedVolumeIds = append(underReplicatedVolumeIds, vid)
		} else if replicaPlacement.GetCopyCount() < len(replicas) {
			overReplicatedVolumeIds = append(overReplicatedVolumeIds, vid)
			fmt.Fprintf(writer, "volume %d replication %s, but over replicated %+d\n", replica.info.Id, replicaPlacement, len(replicas))
		}
	}

	if len(overReplicatedVolumeIds) > 0 {
		return c.fixOverReplicatedVolumes(commandEnv, writer, takeAction, overReplicatedVolumeIds, volumeReplicas, allLocations)
	}

	if len(underReplicatedVolumeIds) == 0 {
		return nil
	}

	// find the most under populated data nodes
	keepDataNodesSorted(allLocations)

	return c.fixUnderReplicatedVolumes(commandEnv, writer, takeAction, underReplicatedVolumeIds, volumeReplicas, allLocations)

}

func collectVolumeReplicaLocations(resp *master_pb.VolumeListResponse) (map[uint32][]*VolumeReplica, []location) {
	volumeReplicas := make(map[uint32][]*VolumeReplica)
	var allLocations []location
	eachDataNode(resp.TopologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		loc := newLocation(dc, string(rack), dn)
		for _, v := range dn.VolumeInfos {
			volumeReplicas[v.Id] = append(volumeReplicas[v.Id], &VolumeReplica{
				location: &loc,
				info:     v,
			})
		}
		allLocations = append(allLocations, loc)
	})
	return volumeReplicas, allLocations
}

func (c *commandVolumeFixReplication) fixOverReplicatedVolumes(commandEnv *CommandEnv, writer io.Writer, takeAction bool, overReplicatedVolumeIds []uint32, volumeReplicas map[uint32][]*VolumeReplica, allLocations []location) error {
	for _, vid := range overReplicatedVolumeIds {
		replicas := volumeReplicas[vid]
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(replicas[0].info.ReplicaPlacement))

		replica := pickOneReplicaToDelete(replicas, replicaPlacement)

		// check collection name pattern
		if *c.collectionPattern != "" {
			matched, err := filepath.Match(*c.collectionPattern, replica.info.Collection)
			if err != nil {
				return fmt.Errorf("match pattern %s with collection %s: %v", *c.collectionPattern, replica.info.Collection, err)
			}
			if !matched {
				break
			}
		}

		fmt.Fprintf(writer, "deleting volume %d from %s ...\n", replica.info.Id, replica.location.dataNode.Id)

		if !takeAction {
			break
		}

		if err := deleteVolume(commandEnv.option.GrpcDialOption, needle.VolumeId(replica.info.Id), replica.location.dataNode.Id); err != nil {
			return fmt.Errorf("deleting volume %d from %s : %v", replica.info.Id, replica.location.dataNode.Id, err)
		}

	}
	return nil
}

func (c *commandVolumeFixReplication) fixUnderReplicatedVolumes(commandEnv *CommandEnv, writer io.Writer, takeAction bool, underReplicatedVolumeIds []uint32, volumeReplicas map[uint32][]*VolumeReplica, allLocations []location) error {
	for _, vid := range underReplicatedVolumeIds {
		replicas := volumeReplicas[vid]
		replica := pickOneReplicaToCopyFrom(replicas)
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(replica.info.ReplicaPlacement))
		foundNewLocation := false
		hasSkippedCollection := false
		for _, dst := range allLocations {
			// check whether data nodes satisfy the constraints
			if dst.dataNode.FreeVolumeCount > 0 && satisfyReplicaPlacement(replicaPlacement, replicas, dst) {
				// check collection name pattern
				if *c.collectionPattern != "" {
					matched, err := filepath.Match(*c.collectionPattern, replica.info.Collection)
					if err != nil {
						return fmt.Errorf("match pattern %s with collection %s: %v", *c.collectionPattern, replica.info.Collection, err)
					}
					if !matched {
						hasSkippedCollection = true
						break
					}
				}

				// ask the volume server to replicate the volume
				foundNewLocation = true
				fmt.Fprintf(writer, "replicating volume %d %s from %s to dataNode %s ...\n", replica.info.Id, replicaPlacement, replica.location.dataNode.Id, dst.dataNode.Id)

				if !takeAction {
					break
				}

				err := operation.WithVolumeServerClient(dst.dataNode.Id, commandEnv.option.GrpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
					_, replicateErr := volumeServerClient.VolumeCopy(context.Background(), &volume_server_pb.VolumeCopyRequest{
						VolumeId:       replica.info.Id,
						SourceDataNode: replica.location.dataNode.Id,
					})
					if replicateErr != nil {
						return fmt.Errorf("copying from %s => %s : %v", replica.location.dataNode.Id, dst.dataNode.Id, replicateErr)
					}
					return nil
				})

				if err != nil {
					return err
				}

				// adjust free volume count
				dst.dataNode.FreeVolumeCount--
				keepDataNodesSorted(allLocations)
				break
			}
		}
		if !foundNewLocation && !hasSkippedCollection {
			fmt.Fprintf(writer, "failed to place volume %d replica as %s, existing:%+v\n", replica.info.Id, replicaPlacement, len(replicas))
		}

	}
	return nil
}

func keepDataNodesSorted(dataNodes []location) {
	sort.Slice(dataNodes, func(i, j int) bool {
		return dataNodes[i].dataNode.FreeVolumeCount > dataNodes[j].dataNode.FreeVolumeCount
	})
}

/*
  if on an existing data node {
    return false
  }
  if different from existing dcs {
    if lack on different dcs {
      return true
    }else{
      return false
    }
  }
  if not on primary dc {
    return false
  }
  if different from existing racks {
    if lack on different racks {
      return true
    }else{
      return false
    }
  }
  if not on primary rack {
    return false
  }
  if lacks on same rack {
    return true
  } else {
    return false
  }
*/
func satisfyReplicaPlacement(replicaPlacement *super_block.ReplicaPlacement, replicas []*VolumeReplica, possibleLocation location) bool {

	existingDataCenters, _, existingDataNodes := countReplicas(replicas)

	if _, found := existingDataNodes[possibleLocation.String()]; found {
		// avoid duplicated volume on the same data node
		return false
	}

	primaryDataCenters, _ := findTopKeys(existingDataCenters)

	// ensure data center count is within limit
	if _, found := existingDataCenters[possibleLocation.DataCenter()]; !found {
		// different from existing dcs
		if len(existingDataCenters) < replicaPlacement.DiffDataCenterCount+1 {
			// lack on different dcs
			return true
		} else {
			// adding this would go over the different dcs limit
			return false
		}
	}
	// now this is same as one of the existing data center
	if !isAmong(possibleLocation.DataCenter(), primaryDataCenters) {
		// not on one of the primary dcs
		return false
	}

	// now this is one of the primary dcs
	primaryDcRacks := make(map[string]int)
	for _, replica := range replicas {
		if replica.location.DataCenter() != possibleLocation.DataCenter() {
			continue
		}
		primaryDcRacks[replica.location.Rack()] += 1
	}
	primaryRacks, _ := findTopKeys(primaryDcRacks)
	sameRackCount := primaryDcRacks[possibleLocation.Rack()]

	// ensure rack count is within limit
	if _, found := primaryDcRacks[possibleLocation.Rack()]; !found {
		// different from existing racks
		if len(primaryDcRacks) < replicaPlacement.DiffRackCount+1 {
			// lack on different racks
			return true
		} else {
			// adding this would go over the different racks limit
			return false
		}
	}
	// now this is same as one of the existing racks
	if !isAmong(possibleLocation.Rack(), primaryRacks) {
		// not on the primary rack
		return false
	}

	// now this is on the primary rack

	// different from existing data nodes
	if sameRackCount < replicaPlacement.SameRackCount+1 {
		// lack on same rack
		return true
	} else {
		// adding this would go over the same data node limit
		return false
	}

}

func findTopKeys(m map[string]int) (topKeys []string, max int) {
	for k, c := range m {
		if max < c {
			topKeys = topKeys[:0]
			topKeys = append(topKeys, k)
			max = c
		} else if max == c {
			topKeys = append(topKeys, k)
		}
	}
	return
}

func isAmong(key string, keys []string) bool {
	for _, k := range keys {
		if k == key {
			return true
		}
	}
	return false
}

type VolumeReplica struct {
	location *location
	info     *master_pb.VolumeInformationMessage
}

type location struct {
	dc       string
	rack     string
	dataNode *master_pb.DataNodeInfo
}

func newLocation(dc, rack string, dataNode *master_pb.DataNodeInfo) location {
	return location{
		dc:       dc,
		rack:     rack,
		dataNode: dataNode,
	}
}

func (l location) String() string {
	return fmt.Sprintf("%s %s %s", l.dc, l.rack, l.dataNode.Id)
}

func (l location) Rack() string {
	return fmt.Sprintf("%s %s", l.dc, l.rack)
}

func (l location) DataCenter() string {
	return l.dc
}

func pickOneReplicaToCopyFrom(replicas []*VolumeReplica) *VolumeReplica {
	mostRecent := replicas[0]
	for _, replica := range replicas {
		if replica.info.ModifiedAtSecond > mostRecent.info.ModifiedAtSecond {
			mostRecent = replica
		}
	}
	return mostRecent
}

func countReplicas(replicas []*VolumeReplica) (diffDc, diffRack, diffNode map[string]int) {
	diffDc = make(map[string]int)
	diffRack = make(map[string]int)
	diffNode = make(map[string]int)
	for _, replica := range replicas {
		diffDc[replica.location.DataCenter()] += 1
		diffRack[replica.location.Rack()] += 1
		diffNode[replica.location.String()] += 1
	}
	return
}

func pickOneReplicaToDelete(replicas []*VolumeReplica, replicaPlacement *super_block.ReplicaPlacement) *VolumeReplica {

	sort.Slice(replicas, func(i, j int) bool {
		a, b := replicas[i], replicas[j]
		if a.info.CompactRevision != b.info.CompactRevision {
			return a.info.CompactRevision < b.info.CompactRevision
		}
		if a.info.ModifiedAtSecond != b.info.ModifiedAtSecond {
			return a.info.ModifiedAtSecond < b.info.ModifiedAtSecond
		}
		if a.info.Size != b.info.Size {
			return a.info.Size < b.info.Size
		}
		return false
	})

	return replicas[0]

}
<|endoftext|>