<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package ssh_test

import (
	"encoding/base64"
	"strings"
	stdtesting "testing"

	gc "launchpad.net/gocheck"

	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/testing/testbase"
	"launchpad.net/juju-core/utils/ssh"
	sshtesting "launchpad.net/juju-core/utils/ssh/testing"
)

func Test(t *stdtesting.T) {
	gc.TestingT(t)
}

type AuthorisedKeysKeysSuite struct {
	testbase.LoggingSuite
}

const (
	// We'll use the current user for ssh tests.
	testSSHUser = ""
)

var _ = gc.Suite(&AuthorisedKeysKeysSuite{})

func (s *AuthorisedKeysKeysSuite) SetUpTest(c *gc.C) {
	s.LoggingSuite.SetUpTest(c)
	fakeHome := coretesting.MakeEmptyFakeHomeWithoutJuju(c)
	s.AddCleanup(func(*gc.C) { fakeHome.Restore() })
}

func writeAuthKeysFile(c *gc.C, keys []string) {
	err := ssh.WriteAuthorisedKeys(testSSHUser, keys)
	c.Assert(err, gc.IsNil)
}

func (s *AuthorisedKeysKeysSuite) TestListKeys(c *gc.C) {
	keys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key,
	}
	writeAuthKeysFile(c, keys)
	keys, err := ssh.ListKeys(testSSHUser, ssh.Fingerprints)
	c.Assert(err, gc.IsNil)
	c.Assert(
		keys, gc.DeepEquals,
		[]string{sshtesting.ValidKeyOne.Fingerprint + " (user@host)", sshtesting.ValidKeyTwo.Fingerprint})
}

func (s *AuthorisedKeysKeysSuite) TestListKeysFull(c *gc.C) {
	keys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key + " anotheruser@host",
	}
	writeAuthKeysFile(c, keys)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, keys)
}

func (s *AuthorisedKeysKeysSuite) TestAddNewKey(c *gc.C) {
	key := sshtesting.ValidKeyOne.Key + " user@host"
	err := ssh.AddKeys(testSSHUser, key)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{key})
}

func (s *AuthorisedKeysKeysSuite) TestAddMoreKeys(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	writeAuthKeysFile(c, []string{firstKey})
	moreKeys := []string{
		sshtesting.ValidKeyTwo.Key + " anotheruser@host",
		sshtesting.ValidKeyThree.Key + " yetanotheruser@host",
	}
	err := ssh.AddKeys(testSSHUser, moreKeys...)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, append([]string{firstKey}, moreKeys...))
}

func (s *AuthorisedKeysKeysSuite) TestAddDuplicateKey(c *gc.C) {
	key := sshtesting.ValidKeyOne.Key + " user@host"
	err := ssh.AddKeys(testSSHUser, key)
	c.Assert(err, gc.IsNil)
	moreKeys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key + " yetanotheruser@host",
	}
	err = ssh.AddKeys(testSSHUser, moreKeys...)
	c.Assert(err, gc.ErrorMatches, "cannot add duplicate ssh key: "+sshtesting.ValidKeyOne.Fingerprint)
}

func (s *AuthorisedKeysKeysSuite) TestAddDuplicateComment(c *gc.C) {
	key := sshtesting.ValidKeyOne.Key + " user@host"
	err := ssh.AddKeys(testSSHUser, key)
	c.Assert(err, gc.IsNil)
	moreKeys := []string{
		sshtesting.ValidKeyTwo.Key + " user@host",
		sshtesting.ValidKeyThree.Key + " yetanotheruser@host",
	}
	err = ssh.AddKeys(testSSHUser, moreKeys...)
	c.Assert(err, gc.ErrorMatches, "cannot add ssh key with duplicate comment: user@host")
}

func (s *AuthorisedKeysKeysSuite) TestAddKeyWithoutComment(c *gc.C) {
	keys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key,
	}
	err := ssh.AddKeys(testSSHUser, keys...)
	c.Assert(err, gc.ErrorMatches, "cannot add ssh key without comment")
}

func (s *AuthorisedKeysKeysSuite) TestAddKeepsUnrecognised(c *gc.C) {
	writeAuthKeysFile(c, []string{sshtesting.ValidKeyOne.Key, "invalid-key"})
	anotherKey := sshtesting.ValidKeyTwo.Key + " anotheruser@host"
	err := ssh.AddKeys(testSSHUser, anotherKey)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ReadAuthorisedKeys(testSSHUser)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{sshtesting.ValidKeyOne.Key, "invalid-key", anotherKey})
}

func (s *AuthorisedKeysKeysSuite) TestDeleteKeys(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	anotherKey := sshtesting.ValidKeyTwo.Key
	thirdKey := sshtesting.ValidKeyThree.Key + " anotheruser@host"
	writeAuthKeysFile(c, []string{firstKey, anotherKey, thirdKey})
	err := ssh.DeleteKeys(testSSHUser, "user@host", sshtesting.ValidKeyTwo.Fingerprint)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{thirdKey})
}

func (s *AuthorisedKeysKeysSuite) TestDeleteKeysKeepsUnrecognised(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	writeAuthKeysFile(c, []string{firstKey, sshtesting.ValidKeyTwo.Key, "invalid-key"})
	err := ssh.DeleteKeys(testSSHUser, "user@host")
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ReadAuthorisedKeys(testSSHUser)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{"invalid-key", sshtesting.ValidKeyTwo.Key})
}

func (s *AuthorisedKeysKeysSuite) TestDeleteNonExistentComment(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	writeAuthKeysFile(c, []string{firstKey})
	err := ssh.DeleteKeys(testSSHUser, "someone@host")
	c.Assert(err, gc.ErrorMatches, "cannot delete non existent key: someone@host")
}

func (s *AuthorisedKeysKeysSuite) TestDeleteNonExistentFingerprint(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	writeAuthKeysFile(c, []string{firstKey})
	err := ssh.DeleteKeys(testSSHUser, sshtesting.ValidKeyTwo.Fingerprint)
	c.Assert(err, gc.ErrorMatches, "cannot delete non existent key: "+sshtesting.ValidKeyTwo.Fingerprint)
}

func (s *AuthorisedKeysKeysSuite) TestDeleteLastKeyForbidden(c *gc.C) {
	keys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key + " yetanotheruser@host",
	}
	writeAuthKeysFile(c, keys)
	err := ssh.DeleteKeys(testSSHUser, "user@host", sshtesting.ValidKeyTwo.Fingerprint)
	c.Assert(err, gc.ErrorMatches, "cannot delete all keys")
}

func (s *AuthorisedKeysKeysSuite) TestReplaceKeys(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	anotherKey := sshtesting.ValidKeyTwo.Key
	writeAuthKeysFile(c, []string{firstKey, anotherKey})

	replaceKey := sshtesting.ValidKeyThree.Key
	err := ssh.ReplaceKeys(testSSHUser, replaceKey)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{replaceKey})
}

func (s *AuthorisedKeysKeysSuite) TestReplaceKeepsUnrecognised(c *gc.C) {
	writeAuthKeysFile(c, []string{sshtesting.ValidKeyOne.Key, "invalid-key"})
	anotherKey := sshtesting.ValidKeyTwo.Key + " anotheruser@host"
	err := ssh.ReplaceKeys(testSSHUser, anotherKey)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ReadAuthorisedKeys(testSSHUser)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{"invalid-key", anotherKey})
}

func (s *AuthorisedKeysKeysSuite) TestEnsureJujuComment(c *gc.C) {
	sshKey := sshtesting.ValidKeyOne.Key
	for _, test := range []struct {
		key      string
		expected string
	}{
		{"invalid-key", "invalid-key"},
		{sshKey, sshKey + " Juju:sshkey"},
		{sshKey + " user@host", sshKey + " Juju:user@host"},
		{sshKey + " Juju:user@host", sshKey + " Juju:user@host"},
		{sshKey + " " + sshKey[3:5], sshKey + " Juju:" + sshKey[3:5]},
	} {
		actual := ssh.EnsureJujuComment(test.key)
		c.Assert(actual, gc.Equals, test.expected)
	}
}

func (s *AuthorisedKeysKeysSuite) TestSplitAuthorisedKeys(c *gc.C) {
	sshKey := sshtesting.ValidKeyOne.Key
	for _, test := range []struct {
		keyData  string
		expected []string
	}{
		{"", nil},
		{sshKey, []string{sshKey}},
		{sshKey + "\n", []string{sshKey}},
		{sshKey + "\n\n", []string{sshKey}},
		{sshKey + "\n#comment\n", []string{sshKey}},
		{sshKey + "\n #comment\n", []string{sshKey}},
		{sshKey + "\ninvalid\n", []string{sshKey, "invalid"}},
	} {
		actual := ssh.SplitAuthorisedKeys(test.keyData)
		c.Assert(actual, gc.DeepEquals, test.expected)
	}
}

func b64decode(c *gc.C, s string) []byte {
	b, err := base64.StdEncoding.DecodeString(s)
	c.Assert(err, gc.IsNil)
	return b
}

func (s *AuthorisedKeysKeysSuite) TestParseAuthorisedKey(c *gc.C) {
	for i, test := range []struct {
		line    string
		key     []byte
		comment string
		err     string
	}{{
		line: sshtesting.ValidKeyOne.Key,
		key:  b64decode(c, strings.Fields(sshtesting.ValidKeyOne.Key)[1]),
	}, {
		line:    sshtesting.ValidKeyOne.Key + " a b c",
		key:     b64decode(c, strings.Fields(sshtesting.ValidKeyOne.Key)[1]),
		comment: "a b c",
	}, {
		line: "ssh-xsa blah",
		err:  "invalid authorized_key \"ssh-xsa blah\"",
	}, {
		// options should be skipped
		line: `no-pty,principals="\"",command="\!" ` + sshtesting.ValidKeyOne.Key,
		key:  b64decode(c, strings.Fields(sshtesting.ValidKeyOne.Key)[1]),
	}, {
		line: "ssh-rsa",
		err:  "invalid authorized_key \"ssh-rsa\"",
	}} {
		c.Logf("test %d: %s", i, test.line)
		ak, err := ssh.ParseAuthorisedKey(test.line)
		if test.err != "" {
			c.Assert(err, gc.ErrorMatches, test.err)
		} else {
			c.Assert(err, gc.IsNil)
			c.Assert(ak, gc.Not(gc.IsNil))
			c.Assert(ak.Key, gc.DeepEquals, test.key)
			c.Assert(ak.Comment, gc.Equals, test.comment)
		}
	}
}
<commit_msg>Add comment<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package ssh_test

import (
	"encoding/base64"
	"strings"
	stdtesting "testing"

	gc "launchpad.net/gocheck"

	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/testing/testbase"
	"launchpad.net/juju-core/utils/ssh"
	sshtesting "launchpad.net/juju-core/utils/ssh/testing"
)

func Test(t *stdtesting.T) {
	gc.TestingT(t)
}

type AuthorisedKeysKeysSuite struct {
	testbase.LoggingSuite
}

const (
	// We'll use the current user for ssh tests.
	testSSHUser = ""
)

var _ = gc.Suite(&AuthorisedKeysKeysSuite{})

func (s *AuthorisedKeysKeysSuite) SetUpTest(c *gc.C) {
	s.LoggingSuite.SetUpTest(c)
	fakeHome := coretesting.MakeEmptyFakeHomeWithoutJuju(c)
	s.AddCleanup(func(*gc.C) { fakeHome.Restore() })
}

func writeAuthKeysFile(c *gc.C, keys []string) {
	err := ssh.WriteAuthorisedKeys(testSSHUser, keys)
	c.Assert(err, gc.IsNil)
}

func (s *AuthorisedKeysKeysSuite) TestListKeys(c *gc.C) {
	keys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key,
	}
	writeAuthKeysFile(c, keys)
	keys, err := ssh.ListKeys(testSSHUser, ssh.Fingerprints)
	c.Assert(err, gc.IsNil)
	c.Assert(
		keys, gc.DeepEquals,
		[]string{sshtesting.ValidKeyOne.Fingerprint + " (user@host)", sshtesting.ValidKeyTwo.Fingerprint})
}

func (s *AuthorisedKeysKeysSuite) TestListKeysFull(c *gc.C) {
	keys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key + " anotheruser@host",
	}
	writeAuthKeysFile(c, keys)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, keys)
}

func (s *AuthorisedKeysKeysSuite) TestAddNewKey(c *gc.C) {
	key := sshtesting.ValidKeyOne.Key + " user@host"
	err := ssh.AddKeys(testSSHUser, key)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{key})
}

func (s *AuthorisedKeysKeysSuite) TestAddMoreKeys(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	writeAuthKeysFile(c, []string{firstKey})
	moreKeys := []string{
		sshtesting.ValidKeyTwo.Key + " anotheruser@host",
		sshtesting.ValidKeyThree.Key + " yetanotheruser@host",
	}
	err := ssh.AddKeys(testSSHUser, moreKeys...)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, append([]string{firstKey}, moreKeys...))
}

func (s *AuthorisedKeysKeysSuite) TestAddDuplicateKey(c *gc.C) {
	key := sshtesting.ValidKeyOne.Key + " user@host"
	err := ssh.AddKeys(testSSHUser, key)
	c.Assert(err, gc.IsNil)
	moreKeys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key + " yetanotheruser@host",
	}
	err = ssh.AddKeys(testSSHUser, moreKeys...)
	c.Assert(err, gc.ErrorMatches, "cannot add duplicate ssh key: "+sshtesting.ValidKeyOne.Fingerprint)
}

func (s *AuthorisedKeysKeysSuite) TestAddDuplicateComment(c *gc.C) {
	key := sshtesting.ValidKeyOne.Key + " user@host"
	err := ssh.AddKeys(testSSHUser, key)
	c.Assert(err, gc.IsNil)
	moreKeys := []string{
		sshtesting.ValidKeyTwo.Key + " user@host",
		sshtesting.ValidKeyThree.Key + " yetanotheruser@host",
	}
	err = ssh.AddKeys(testSSHUser, moreKeys...)
	c.Assert(err, gc.ErrorMatches, "cannot add ssh key with duplicate comment: user@host")
}

func (s *AuthorisedKeysKeysSuite) TestAddKeyWithoutComment(c *gc.C) {
	keys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key,
	}
	err := ssh.AddKeys(testSSHUser, keys...)
	c.Assert(err, gc.ErrorMatches, "cannot add ssh key without comment")
}

func (s *AuthorisedKeysKeysSuite) TestAddKeepsUnrecognised(c *gc.C) {
	writeAuthKeysFile(c, []string{sshtesting.ValidKeyOne.Key, "invalid-key"})
	anotherKey := sshtesting.ValidKeyTwo.Key + " anotheruser@host"
	err := ssh.AddKeys(testSSHUser, anotherKey)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ReadAuthorisedKeys(testSSHUser)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{sshtesting.ValidKeyOne.Key, "invalid-key", anotherKey})
}

func (s *AuthorisedKeysKeysSuite) TestDeleteKeys(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	anotherKey := sshtesting.ValidKeyTwo.Key
	thirdKey := sshtesting.ValidKeyThree.Key + " anotheruser@host"
	writeAuthKeysFile(c, []string{firstKey, anotherKey, thirdKey})
	err := ssh.DeleteKeys(testSSHUser, "user@host", sshtesting.ValidKeyTwo.Fingerprint)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{thirdKey})
}

func (s *AuthorisedKeysKeysSuite) TestDeleteKeysKeepsUnrecognised(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	writeAuthKeysFile(c, []string{firstKey, sshtesting.ValidKeyTwo.Key, "invalid-key"})
	err := ssh.DeleteKeys(testSSHUser, "user@host")
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ReadAuthorisedKeys(testSSHUser)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{"invalid-key", sshtesting.ValidKeyTwo.Key})
}

func (s *AuthorisedKeysKeysSuite) TestDeleteNonExistentComment(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	writeAuthKeysFile(c, []string{firstKey})
	err := ssh.DeleteKeys(testSSHUser, "someone@host")
	c.Assert(err, gc.ErrorMatches, "cannot delete non existent key: someone@host")
}

func (s *AuthorisedKeysKeysSuite) TestDeleteNonExistentFingerprint(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	writeAuthKeysFile(c, []string{firstKey})
	err := ssh.DeleteKeys(testSSHUser, sshtesting.ValidKeyTwo.Fingerprint)
	c.Assert(err, gc.ErrorMatches, "cannot delete non existent key: "+sshtesting.ValidKeyTwo.Fingerprint)
}

func (s *AuthorisedKeysKeysSuite) TestDeleteLastKeyForbidden(c *gc.C) {
	keys := []string{
		sshtesting.ValidKeyOne.Key + " user@host",
		sshtesting.ValidKeyTwo.Key + " yetanotheruser@host",
	}
	writeAuthKeysFile(c, keys)
	err := ssh.DeleteKeys(testSSHUser, "user@host", sshtesting.ValidKeyTwo.Fingerprint)
	c.Assert(err, gc.ErrorMatches, "cannot delete all keys")
}

func (s *AuthorisedKeysKeysSuite) TestReplaceKeys(c *gc.C) {
	firstKey := sshtesting.ValidKeyOne.Key + " user@host"
	anotherKey := sshtesting.ValidKeyTwo.Key
	writeAuthKeysFile(c, []string{firstKey, anotherKey})

	// replaceKey is created without a comment so test that
	// ReplaceKeys handles keys without comments. This is
	// because existing keys may not have a comment and
	// ReplaceKeys is used to rewrite the entire authorized_keys
	// file when adding new keys.
	replaceKey := sshtesting.ValidKeyThree.Key
	err := ssh.ReplaceKeys(testSSHUser, replaceKey)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ListKeys(testSSHUser, ssh.FullKeys)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{replaceKey})
}

func (s *AuthorisedKeysKeysSuite) TestReplaceKeepsUnrecognised(c *gc.C) {
	writeAuthKeysFile(c, []string{sshtesting.ValidKeyOne.Key, "invalid-key"})
	anotherKey := sshtesting.ValidKeyTwo.Key + " anotheruser@host"
	err := ssh.ReplaceKeys(testSSHUser, anotherKey)
	c.Assert(err, gc.IsNil)
	actual, err := ssh.ReadAuthorisedKeys(testSSHUser)
	c.Assert(err, gc.IsNil)
	c.Assert(actual, gc.DeepEquals, []string{"invalid-key", anotherKey})
}

func (s *AuthorisedKeysKeysSuite) TestEnsureJujuComment(c *gc.C) {
	sshKey := sshtesting.ValidKeyOne.Key
	for _, test := range []struct {
		key      string
		expected string
	}{
		{"invalid-key", "invalid-key"},
		{sshKey, sshKey + " Juju:sshkey"},
		{sshKey + " user@host", sshKey + " Juju:user@host"},
		{sshKey + " Juju:user@host", sshKey + " Juju:user@host"},
		{sshKey + " " + sshKey[3:5], sshKey + " Juju:" + sshKey[3:5]},
	} {
		actual := ssh.EnsureJujuComment(test.key)
		c.Assert(actual, gc.Equals, test.expected)
	}
}

func (s *AuthorisedKeysKeysSuite) TestSplitAuthorisedKeys(c *gc.C) {
	sshKey := sshtesting.ValidKeyOne.Key
	for _, test := range []struct {
		keyData  string
		expected []string
	}{
		{"", nil},
		{sshKey, []string{sshKey}},
		{sshKey + "\n", []string{sshKey}},
		{sshKey + "\n\n", []string{sshKey}},
		{sshKey + "\n#comment\n", []string{sshKey}},
		{sshKey + "\n #comment\n", []string{sshKey}},
		{sshKey + "\ninvalid\n", []string{sshKey, "invalid"}},
	} {
		actual := ssh.SplitAuthorisedKeys(test.keyData)
		c.Assert(actual, gc.DeepEquals, test.expected)
	}
}

func b64decode(c *gc.C, s string) []byte {
	b, err := base64.StdEncoding.DecodeString(s)
	c.Assert(err, gc.IsNil)
	return b
}

func (s *AuthorisedKeysKeysSuite) TestParseAuthorisedKey(c *gc.C) {
	for i, test := range []struct {
		line    string
		key     []byte
		comment string
		err     string
	}{{
		line: sshtesting.ValidKeyOne.Key,
		key:  b64decode(c, strings.Fields(sshtesting.ValidKeyOne.Key)[1]),
	}, {
		line:    sshtesting.ValidKeyOne.Key + " a b c",
		key:     b64decode(c, strings.Fields(sshtesting.ValidKeyOne.Key)[1]),
		comment: "a b c",
	}, {
		line: "ssh-xsa blah",
		err:  "invalid authorized_key \"ssh-xsa blah\"",
	}, {
		// options should be skipped
		line: `no-pty,principals="\"",command="\!" ` + sshtesting.ValidKeyOne.Key,
		key:  b64decode(c, strings.Fields(sshtesting.ValidKeyOne.Key)[1]),
	}, {
		line: "ssh-rsa",
		err:  "invalid authorized_key \"ssh-rsa\"",
	}} {
		c.Logf("test %d: %s", i, test.line)
		ak, err := ssh.ParseAuthorisedKey(test.line)
		if test.err != "" {
			c.Assert(err, gc.ErrorMatches, test.err)
		} else {
			c.Assert(err, gc.IsNil)
			c.Assert(ak, gc.Not(gc.IsNil))
			c.Assert(ak.Key, gc.DeepEquals, test.key)
			c.Assert(ak.Comment, gc.Equals, test.comment)
		}
	}
}
<|endoftext|>


<commit_before>package checkawscloudwatchlogs

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs/cloudwatchlogsiface"
	"github.com/jessevdk/go-flags"

	"github.com/mackerelio/checkers"
	"github.com/mackerelio/golib/pluginutil"
	"github.com/natefinch/atomic"
)

type logOpts struct {
	LogGroupName        string `long:"log-group-name" required:"true" value-name:"LOG-GROUP-NAME" description:"Log group name" unquote:"false"`
	LogStreamNamePrefix string `long:"log-stream-name-prefix" value-name:"LOG-STREAM-NAME-PREFIX" description:"Log stream name prefix" unquote:"false"`

	Pattern       string `short:"p" long:"pattern" required:"true" value-name:"PATTERN" description:"Pattern to search for. The value is recognized as the pattern syntax of CloudWatch Logs." unquote:"false"`
	WarningOver   int    `short:"w" long:"warning-over" value-name:"WARNING" description:"Trigger a warning if matched lines is over a number"`
	CriticalOver  int    `short:"c" long:"critical-over" value-name:"CRITICAL" description:"Trigger a critical if matched lines is over a number"`
	StateDir      string `short:"s" long:"state-dir" value-name:"DIR" description:"Dir to keep state files under" unquote:"false"`
	ReturnContent bool   `short:"r" long:"return" description:"Output matched lines"`
	MaxRetries    int    `short:"t" long:"max-retries" value-name:"MAX-RETRIES" description:"Maximum number of retries to call the AWS API"`
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "CloudWatch Logs"
	ckr.Exit()
}

type awsCloudwatchLogsPlugin struct {
	Service   cloudwatchlogsiface.CloudWatchLogsAPI
	StateFile string
	*logOpts
}

func newCloudwatchLogsPlugin(opts *logOpts, args []string) (*awsCloudwatchLogsPlugin, error) {
	var err error
	p := &awsCloudwatchLogsPlugin{logOpts: opts}
	p.Service, err = createService(opts)
	if err != nil {
		return nil, err
	}
	if p.StateDir == "" {
		workdir := pluginutil.PluginWorkDir()
		p.StateDir = filepath.Join(workdir, "check-cloudwatch-logs")
	}
	p.StateFile = getStateFile(p.StateDir, opts.LogGroupName, opts.LogStreamNamePrefix, args)
	return p, nil
}

var stateRe = regexp.MustCompile(`[^-a-zA-Z0-9_.]`)

func getStateFile(stateDir, logGroupName, logStreamNamePrefix string, args []string) string {
	return filepath.Join(
		stateDir,
		fmt.Sprintf(
			"%s-%x.json",
			strings.TrimLeft(stateRe.ReplaceAllString(logGroupName+"_"+logStreamNamePrefix, "_"), "_"),
			md5.Sum([]byte(
				strings.Join(
					[]string{
						os.Getenv("AWS_PROFILE"),
						os.Getenv("AWS_ACCESS_KEY_ID"),
						os.Getenv("AWS_SECRET_ACCESS_KEY"),
						os.Getenv("AWS_REGION"),
						strings.Join(args, " "),
					},
					" ",
				)),
			),
		),
	)
}

func createAWSConfig(opts *logOpts) *aws.Config {
	conf := aws.NewConfig()
	if opts.MaxRetries > 0 {
		return conf.WithMaxRetries(opts.MaxRetries)
	}
	return conf
}

func createService(opts *logOpts) (*cloudwatchlogs.CloudWatchLogs, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}
	return cloudwatchlogs.New(sess, createAWSConfig(opts)), nil
}

type logState struct {
	NextToken *string
	StartTime *int64
}

func (p *awsCloudwatchLogsPlugin) collect(now time.Time) ([]string, error) {
	var nextToken *string
	var startTime *int64
	if s, err := p.loadState(); err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		}
	} else {
		if s.StartTime != nil && *s.StartTime > now.Add(-time.Hour).Unix()*1000 {
			nextToken = s.NextToken
			startTime = s.StartTime
		}
	}
	if startTime == nil {
		startTime = aws.Int64(now.Add(-1*time.Minute).Unix() * 1000)
	}
	var messages []string
	input := &cloudwatchlogs.FilterLogEventsInput{
		StartTime:     startTime,
		LogGroupName:  aws.String(p.LogGroupName),
		NextToken:     nextToken,
		FilterPattern: aws.String(p.Pattern),
	}
	if p.LogStreamNamePrefix != "" {
		input.LogStreamNamePrefix = aws.String(p.LogStreamNamePrefix)
	}
	err := p.Service.FilterLogEventsPages(input, func(output *cloudwatchlogs.FilterLogEventsOutput, lastPage bool) bool {
		for _, event := range output.Events {
			messages = append(messages, *event.Message)
			if startTime == nil || *startTime <= *event.Timestamp {
				startTime = aws.Int64(*event.Timestamp + 1)
			}
		}
		nextToken = output.NextToken
		time.Sleep(150 * time.Millisecond)
		return true
	})
	if err != nil {
		return nil, err
	}
	if err := p.saveState(&logState{nextToken, startTime}); err != nil {
		return nil, err
	}
	return messages, nil
}

func (p *awsCloudwatchLogsPlugin) loadState() (*logState, error) {
	f, err := os.Open(p.StateFile)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	var s logState
	err = json.NewDecoder(f).Decode(&s)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func (p *awsCloudwatchLogsPlugin) saveState(s *logState) error {
	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(s); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(p.StateFile), 0755); err != nil {
		return err
	}
	return atomic.WriteFile(p.StateFile, &buf)
}

func (p *awsCloudwatchLogsPlugin) check(messages []string) *checkers.Checker {
	status := checkers.OK
	msg := fmt.Sprint(len(messages))
	if len(messages) > p.CriticalOver {
		status = checkers.CRITICAL
		msg += " > " + fmt.Sprint(p.CriticalOver)
	} else if len(messages) > p.WarningOver {
		status = checkers.WARNING
		msg += " > " + fmt.Sprint(p.WarningOver)
	}
	msg += " messages for pattern /" + p.Pattern + "/"
	if status != checkers.OK && p.ReturnContent {
		msg += "\n" + strings.Join(messages, "")
	}
	return checkers.NewChecker(status, msg)
}

func (p *awsCloudwatchLogsPlugin) run(now time.Time) *checkers.Checker {
	messages, err := p.collect(now)
	if err != nil {
		return checkers.Unknown(fmt.Sprint(err))
	}
	return p.check(messages)
}

func run(args []string) *checkers.Checker {
	opts := &logOpts{}
	_, err := flags.ParseArgs(opts, args)
	if err != nil {
		os.Exit(1)
	}
	p, err := newCloudwatchLogsPlugin(opts, args)
	if err != nil {
		return checkers.Unknown(fmt.Sprint(err))
	}
	return p.run(time.Now())
}
<commit_msg>refactor collect<commit_after>package checkawscloudwatchlogs

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs/cloudwatchlogsiface"
	"github.com/jessevdk/go-flags"

	"github.com/mackerelio/checkers"
	"github.com/mackerelio/golib/pluginutil"
	"github.com/natefinch/atomic"
)

type logOpts struct {
	LogGroupName        string `long:"log-group-name" required:"true" value-name:"LOG-GROUP-NAME" description:"Log group name" unquote:"false"`
	LogStreamNamePrefix string `long:"log-stream-name-prefix" value-name:"LOG-STREAM-NAME-PREFIX" description:"Log stream name prefix" unquote:"false"`

	Pattern       string `short:"p" long:"pattern" required:"true" value-name:"PATTERN" description:"Pattern to search for. The value is recognized as the pattern syntax of CloudWatch Logs." unquote:"false"`
	WarningOver   int    `short:"w" long:"warning-over" value-name:"WARNING" description:"Trigger a warning if matched lines is over a number"`
	CriticalOver  int    `short:"c" long:"critical-over" value-name:"CRITICAL" description:"Trigger a critical if matched lines is over a number"`
	StateDir      string `short:"s" long:"state-dir" value-name:"DIR" description:"Dir to keep state files under" unquote:"false"`
	ReturnContent bool   `short:"r" long:"return" description:"Output matched lines"`
	MaxRetries    int    `short:"t" long:"max-retries" value-name:"MAX-RETRIES" description:"Maximum number of retries to call the AWS API"`
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "CloudWatch Logs"
	ckr.Exit()
}

type awsCloudwatchLogsPlugin struct {
	Service   cloudwatchlogsiface.CloudWatchLogsAPI
	StateFile string
	*logOpts
}

func newCloudwatchLogsPlugin(opts *logOpts, args []string) (*awsCloudwatchLogsPlugin, error) {
	var err error
	p := &awsCloudwatchLogsPlugin{logOpts: opts}
	p.Service, err = createService(opts)
	if err != nil {
		return nil, err
	}
	if p.StateDir == "" {
		workdir := pluginutil.PluginWorkDir()
		p.StateDir = filepath.Join(workdir, "check-cloudwatch-logs")
	}
	p.StateFile = getStateFile(p.StateDir, opts.LogGroupName, opts.LogStreamNamePrefix, args)
	return p, nil
}

var stateRe = regexp.MustCompile(`[^-a-zA-Z0-9_.]`)

func getStateFile(stateDir, logGroupName, logStreamNamePrefix string, args []string) string {
	return filepath.Join(
		stateDir,
		fmt.Sprintf(
			"%s-%x.json",
			strings.TrimLeft(stateRe.ReplaceAllString(logGroupName+"_"+logStreamNamePrefix, "_"), "_"),
			md5.Sum([]byte(
				strings.Join(
					[]string{
						os.Getenv("AWS_PROFILE"),
						os.Getenv("AWS_ACCESS_KEY_ID"),
						os.Getenv("AWS_SECRET_ACCESS_KEY"),
						os.Getenv("AWS_REGION"),
						strings.Join(args, " "),
					},
					" ",
				)),
			),
		),
	)
}

func createAWSConfig(opts *logOpts) *aws.Config {
	conf := aws.NewConfig()
	if opts.MaxRetries > 0 {
		return conf.WithMaxRetries(opts.MaxRetries)
	}
	return conf
}

func createService(opts *logOpts) (*cloudwatchlogs.CloudWatchLogs, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}
	return cloudwatchlogs.New(sess, createAWSConfig(opts)), nil
}

type logState struct {
	NextToken *string
	StartTime *int64
}

func (p *awsCloudwatchLogsPlugin) collect(now time.Time) ([]string, error) {
	s, err := p.loadState()
	if err != nil {
		return nil, err
	}
	if s.StartTime == nil || *s.StartTime <= now.Add(-time.Hour).Unix()*1000 {
		s.StartTime = aws.Int64(now.Add(-1*time.Minute).Unix() * 1000)
	}
	var messages []string
	input := &cloudwatchlogs.FilterLogEventsInput{
		StartTime:     s.StartTime,
		LogGroupName:  aws.String(p.LogGroupName),
		NextToken:     s.NextToken,
		FilterPattern: aws.String(p.Pattern),
	}
	if p.LogStreamNamePrefix != "" {
		input.LogStreamNamePrefix = aws.String(p.LogStreamNamePrefix)
	}
	err = p.Service.FilterLogEventsPages(input, func(output *cloudwatchlogs.FilterLogEventsOutput, lastPage bool) bool {
		for _, event := range output.Events {
			messages = append(messages, *event.Message)
			if s.StartTime == nil || *s.StartTime <= *event.Timestamp {
				s.StartTime = aws.Int64(*event.Timestamp + 1)
			}
		}
		s.NextToken = output.NextToken
		time.Sleep(150 * time.Millisecond)
		return true
	})
	if err != nil {
		return nil, err
	}
	err = p.saveState(s)
	if err != nil {
		return nil, err
	}
	return messages, nil
}

func (p *awsCloudwatchLogsPlugin) loadState() (*logState, error) {
	var s logState
	f, err := os.Open(p.StateFile)
	defer f.Close()
	if err != nil {
		if os.IsNotExist(err) {
			return &s, nil
		}
		return nil, err
	}
	err = json.NewDecoder(f).Decode(&s)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func (p *awsCloudwatchLogsPlugin) saveState(s *logState) error {
	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(s); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(p.StateFile), 0755); err != nil {
		return err
	}
	return atomic.WriteFile(p.StateFile, &buf)
}

func (p *awsCloudwatchLogsPlugin) check(messages []string) *checkers.Checker {
	status := checkers.OK
	msg := fmt.Sprint(len(messages))
	if len(messages) > p.CriticalOver {
		status = checkers.CRITICAL
		msg += " > " + fmt.Sprint(p.CriticalOver)
	} else if len(messages) > p.WarningOver {
		status = checkers.WARNING
		msg += " > " + fmt.Sprint(p.WarningOver)
	}
	msg += " messages for pattern /" + p.Pattern + "/"
	if status != checkers.OK && p.ReturnContent {
		msg += "\n" + strings.Join(messages, "")
	}
	return checkers.NewChecker(status, msg)
}

func (p *awsCloudwatchLogsPlugin) run(now time.Time) *checkers.Checker {
	messages, err := p.collect(now)
	if err != nil {
		return checkers.Unknown(fmt.Sprint(err))
	}
	return p.check(messages)
}

func run(args []string) *checkers.Checker {
	opts := &logOpts{}
	_, err := flags.ParseArgs(opts, args)
	if err != nil {
		os.Exit(1)
	}
	p, err := newCloudwatchLogsPlugin(opts, args)
	if err != nil {
		return checkers.Unknown(fmt.Sprint(err))
	}
	return p.run(time.Now())
}
<|endoftext|>


<commit_before>package pixy

import (
	"fmt"
	"math"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/mailgun/kafka-pixy/Godeps/_workspace/src/github.com/Shopify/sarama"
	. "github.com/mailgun/kafka-pixy/Godeps/_workspace/src/gopkg.in/check.v1"
)

const (
	testSocket = "kafka-pixy.sock"
)

type ServiceSuite struct {
	serviceCfg *ServiceCfg
	tkc        *TestKafkaClient
	unixClient *http.Client
	tcpClient  *http.Client
}

var _ = Suite(&ServiceSuite{})

func (s *ServiceSuite) SetUpSuite(c *C) {
	InitTestLog()
}

func (s *ServiceSuite) SetUpTest(c *C) {
	s.serviceCfg = &ServiceCfg{
		UnixAddr:    path.Join(os.TempDir(), testSocket),
		BrokerAddrs: testKafkaPeers,
	}
	os.Remove(s.serviceCfg.UnixAddr)

	s.tkc = NewTestKafkaClient(s.serviceCfg.BrokerAddrs)
	s.unixClient = NewUDSHTTPClient(s.serviceCfg.UnixAddr)
	s.tcpClient = &http.Client{}
}

func (s *ServiceSuite) TestStartAndStop(c *C) {
	svc, err := SpawnService(s.serviceCfg)
	c.Assert(err, IsNil)
	svc.Stop()
	svc.Wait4Stop()
}

func (s *ServiceSuite) TestInvalidUnixAddr(c *C) {
	// Given
	s.serviceCfg.UnixAddr = "/tmp"
	// When
	svc, err := SpawnService(s.serviceCfg)
	// Then
	c.Assert(err.Error(), Equals,
		"failed to start Unix socket based HTTP API, cause=(failed to create listener, cause=(listen unix /tmp: bind: address already in use))")
	c.Assert(svc, IsNil)
}

func (s *ServiceSuite) TestInvalidBrokers(c *C) {
	// Given
	s.serviceCfg.BrokerAddrs = []string{"localhost:12345"}
	// When
	svc, err := SpawnService(s.serviceCfg)
	// Then
	c.Assert(err.Error(), Equals,
		"failed to spawn Kafka client, cause=(failed to create sarama client, cause=(kafka: Client has run out of available brokers to talk to. Is your cluster reachable?))")
	c.Assert(svc, IsNil)
}

// If `key` is not `nil` then produced messages are deterministically
// distributed between partitions based on the `key` hash.
func (s *ServiceSuite) TestProduce(c *C) {
	// Given
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	for i := 0; i < 10; i++ {
		s.unixClient.Post("http://_/topics/test.4?key=1",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
		s.unixClient.Post("http://_/topics/test.4?key=2",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
		s.unixClient.Post("http://_/topics/test.4?key=3",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
		s.unixClient.Post("http://_/topics/test.4?key=4",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
		s.unixClient.Post("http://_/topics/test.4?key=5",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
	}
	svc.Stop()
	svc.Wait4Stop()
	offsetsAfter := s.tkc.getOffsets("test.4")
	// Then
	c.Assert(offsetsAfter[0], Equals, offsetsBefore[0]+20)
	c.Assert(offsetsAfter[1], Equals, offsetsBefore[1]+10)
	c.Assert(offsetsAfter[2], Equals, offsetsBefore[2]+10)
	c.Assert(offsetsAfter[3], Equals, offsetsBefore[3]+10)
}

// If `key` of a produced message is `nil` then it is submitted to a random
// partition. Therefore a batch of such messages is evenly distributed among
// all available partitions.
func (s *ServiceSuite) TestProduceNilKey(c *C) {
	// Given
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	for i := 0; i < 100; i++ {
		s.unixClient.Post("http://_/topics/test.4",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
	}
	svc.Stop()
	svc.Wait4Stop()
	offsetsAfter := s.tkc.getOffsets("test.4")
	// Then
	delta0 := offsetsAfter[0] - offsetsBefore[0]
	delta1 := offsetsAfter[1] - offsetsBefore[1]
	imbalance := int(math.Abs(float64(delta1 - delta0)))
	if imbalance > 20 {
		panic(fmt.Errorf("Too high imbalance: %v != %v", delta0, delta1))
	}
}

// If `key` of a produced message is empty then it is deterministically
// submitted to a particular partition determined by the empty key hash.
func (s *ServiceSuite) TestProduceEmptyKey(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	for i := 0; i < 10; i++ {
		s.unixClient.Post("http://_/topics/test.4?key=",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
	}
	svc.Stop()
	svc.Wait4Stop()
	offsetsAfter := s.tkc.getOffsets("test.4")
	// Then
	c.Assert(offsetsAfter[0], Equals, offsetsBefore[0])
	c.Assert(offsetsAfter[1], Equals, offsetsBefore[1])
	c.Assert(offsetsAfter[2], Equals, offsetsBefore[2])
	c.Assert(offsetsAfter[3], Equals, offsetsBefore[3]+10)
}

// Utf8 messages are submitted without a problem.
func (s *ServiceSuite) TestUtf8Message(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	s.unixClient.Post("http://_/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Превед Медвед"))
	svc.Stop()
	svc.Wait4Stop()
	// Then
	offsetsAfter := s.tkc.getOffsets("test.4")
	msgs := s.tkc.getMessages("test.4", offsetsBefore, offsetsAfter)
	c.Assert(msgs, DeepEquals,
		[][]string{[]string(nil), []string{"Превед Медвед"}, []string(nil), []string(nil)})
}

// TCP API is not started by default.
func (s *ServiceSuite) TestTCPDoesNotWork(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	// When
	r, err := s.tcpClient.Post("http://localhost:55501/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Hello Kitty"))
	// Then
	svc.Stop()
	svc.Wait4Stop()
	c.Assert(err.Error(), Equals,
		"Post http://localhost:55501/topics/test.4?key=foo: dial tcp 127.0.0.1:55501: connection refused")
	c.Assert(r, IsNil)
}

// API is served on a TCP socket if it is explicitly configured.
func (s *ServiceSuite) TestBothAPI(c *C) {
	offsetsBefore := s.tkc.getOffsets("test.4")
	s.serviceCfg.TCPAddr = "127.0.0.1:55502"
	svc, _ := SpawnService(s.serviceCfg)
	// When
	_, err1 := s.tcpClient.Post("http://localhost:55502/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Превед"))
	_, err2 := s.unixClient.Post("http://_/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Kitty"))
	// Then
	svc.Stop()
	svc.Wait4Stop()
	c.Assert(err1, IsNil)
	c.Assert(err2, IsNil)
	offsetsAfter := s.tkc.getOffsets("test.4")
	msgs := s.tkc.getMessages("test.4", offsetsBefore, offsetsAfter)
	c.Assert(msgs, DeepEquals,
		[][]string{[]string(nil), []string{"Превед", "Kitty"}, []string(nil), []string(nil)})
}

func (s *ServiceSuite) TestStoppedServerCall(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	_, err := s.unixClient.Post("http://_/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Hello"))
	c.Assert(err, IsNil)
	// When
	svc.Stop()
	svc.Wait4Stop()
	// Then
	r, err := s.unixClient.Post("http://_/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Kitty"))
	c.Assert(err.Error(), Equals, "Post http://_/topics/test.4?key=foo: EOF")
	c.Assert(r, IsNil)
}

// If the TCP API Server crashes then the service terminates gracefully.
func (s *ServiceSuite) TestTCPServerCrash(c *C) {
	s.serviceCfg.TCPAddr = "127.0.0.1:55502"
	svc, _ := SpawnService(s.serviceCfg)
	// When
	svc.tcpServer.errorCh <- fmt.Errorf("Kaboom!")
	// Then
	svc.Stop()
	svc.Wait4Stop()
}

// Messages that have maximum possible size indeed go through. Note that we
// assume that the broker's limit is the same as the producer's one or higher.
func (s *ServiceSuite) TestLargestMessage(c *C) {
	offsetsBefore := s.tkc.getOffsets("test.4")
	maxMsgSize := sarama.NewConfig().Producer.MaxMessageBytes - ProdMsgMetadataSize([]byte("foo"))
	msg := GenMessage(maxMsgSize)
	s.serviceCfg.TCPAddr = "127.0.0.1:55503"
	svc, _ := SpawnService(s.serviceCfg)
	// When
	r := PostChunked(s.tcpClient, "http://127.0.0.1:55503/topics/test.4?key=foo", msg)
	svc.Stop()
	svc.Wait4Stop()
	// Then
	AssertHTTPResp(c, r, http.StatusOK, "")
	offsetsAfter := s.tkc.getOffsets("test.4")
	messages := s.tkc.getMessages("test.4", offsetsBefore, offsetsAfter)
	readMsg := messages[1][0]
	c.Assert(readMsg, Equals, msg)
}

// Messages that are larger then producer's MaxMessageBytes size are silently
// dropped. Note that we assume that the broker's limit is the same as the
// producer's one or higher.
func (s *ServiceSuite) TestMessageTooLarge(c *C) {
	offsetsBefore := s.tkc.getOffsets("test.4")
	maxMsgSize := sarama.NewConfig().Producer.MaxMessageBytes - ProdMsgMetadataSize([]byte("foo")) + 1
	msg := GenMessage(maxMsgSize)
	s.serviceCfg.TCPAddr = "127.0.0.1:55504"
	svc, _ := SpawnService(s.serviceCfg)
	// When
	r := PostChunked(s.tcpClient, "http://127.0.0.1:55504/topics/test.4?key=foo", msg)
	svc.Stop()
	svc.Wait4Stop()
	// Then
	AssertHTTPResp(c, r, http.StatusOK, "")
	offsetsAfter := s.tkc.getOffsets("test.4")
	c.Assert(offsetsAfter, DeepEquals, offsetsBefore)
}

func (s *ServiceSuite) TestSyncProduce(c *C) {
	// Given
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	r, err := s.unixClient.Post("http://_/topics/test.4?key=1&sync",
		"text/plain", strings.NewReader("Foo"))
	svc.Stop()
	svc.Wait4Stop()
	offsetsAfter := s.tkc.getOffsets("test.4")
	// Then
	c.Assert(err, IsNil)
	AssertHTTPResp(c, r, http.StatusOK, "")
	c.Assert(offsetsAfter[0], Equals, offsetsBefore[0]+1)
}

func (s *ServiceSuite) TestSyncProduceInvalidTopic(c *C) {
	// Given
	svc, _ := SpawnService(s.serviceCfg)
	// When
	r, err := s.unixClient.Post("http://_/topics/no-such-topic?sync=true",
		"text/plain", strings.NewReader("Foo"))
	svc.Stop()
	svc.Wait4Stop()
	// Then
	c.Assert(err, IsNil)
	AssertHTTPResp(c, r, http.StatusNotFound,
		fmt.Sprintf("%s\n", sarama.ErrUnknownTopicOrPartition.Error()))
}

// If the Unix API Server crashes then the service terminates gracefully.
func (s *ServiceSuite) TestUnixServerCrash(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	// When
	svc.unixServer.errorCh <- fmt.Errorf("Kaboom!")
	// Then
	svc.Stop()
	svc.Wait4Stop()
}
<commit_msg>Fix unit tests failing under golang tip<commit_after>package pixy

import (
	"fmt"
	"math"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/mailgun/kafka-pixy/Godeps/_workspace/src/github.com/Shopify/sarama"
	. "github.com/mailgun/kafka-pixy/Godeps/_workspace/src/gopkg.in/check.v1"
)

const (
	testSocket = "kafka-pixy.sock"
)

type ServiceSuite struct {
	serviceCfg *ServiceCfg
	tkc        *TestKafkaClient
	unixClient *http.Client
	tcpClient  *http.Client
}

var _ = Suite(&ServiceSuite{})

func (s *ServiceSuite) SetUpSuite(c *C) {
	InitTestLog()
}

func (s *ServiceSuite) SetUpTest(c *C) {
	s.serviceCfg = &ServiceCfg{
		UnixAddr:    path.Join(os.TempDir(), testSocket),
		BrokerAddrs: testKafkaPeers,
	}
	os.Remove(s.serviceCfg.UnixAddr)

	s.tkc = NewTestKafkaClient(s.serviceCfg.BrokerAddrs)
	s.unixClient = NewUDSHTTPClient(s.serviceCfg.UnixAddr)
	s.tcpClient = &http.Client{}
}

func (s *ServiceSuite) TestStartAndStop(c *C) {
	svc, err := SpawnService(s.serviceCfg)
	c.Assert(err, IsNil)
	svc.Stop()
	svc.Wait4Stop()
}

func (s *ServiceSuite) TestInvalidUnixAddr(c *C) {
	// Given
	s.serviceCfg.UnixAddr = "/tmp"
	// When
	svc, err := SpawnService(s.serviceCfg)
	// Then
	c.Assert(err.Error(), Equals,
		"failed to start Unix socket based HTTP API, cause=(failed to create listener, cause=(listen unix /tmp: bind: address already in use))")
	c.Assert(svc, IsNil)
}

func (s *ServiceSuite) TestInvalidBrokers(c *C) {
	// Given
	s.serviceCfg.BrokerAddrs = []string{"localhost:12345"}
	// When
	svc, err := SpawnService(s.serviceCfg)
	// Then
	c.Assert(err.Error(), Equals,
		"failed to spawn Kafka client, cause=(failed to create sarama client, cause=(kafka: Client has run out of available brokers to talk to. Is your cluster reachable?))")
	c.Assert(svc, IsNil)
}

// If `key` is not `nil` then produced messages are deterministically
// distributed between partitions based on the `key` hash.
func (s *ServiceSuite) TestProduce(c *C) {
	// Given
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	for i := 0; i < 10; i++ {
		s.unixClient.Post("http://_/topics/test.4?key=1",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
		s.unixClient.Post("http://_/topics/test.4?key=2",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
		s.unixClient.Post("http://_/topics/test.4?key=3",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
		s.unixClient.Post("http://_/topics/test.4?key=4",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
		s.unixClient.Post("http://_/topics/test.4?key=5",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
	}
	svc.Stop()
	svc.Wait4Stop()
	offsetsAfter := s.tkc.getOffsets("test.4")
	// Then
	c.Assert(offsetsAfter[0], Equals, offsetsBefore[0]+20)
	c.Assert(offsetsAfter[1], Equals, offsetsBefore[1]+10)
	c.Assert(offsetsAfter[2], Equals, offsetsBefore[2]+10)
	c.Assert(offsetsAfter[3], Equals, offsetsBefore[3]+10)
}

// If `key` of a produced message is `nil` then it is submitted to a random
// partition. Therefore a batch of such messages is evenly distributed among
// all available partitions.
func (s *ServiceSuite) TestProduceNilKey(c *C) {
	// Given
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	for i := 0; i < 100; i++ {
		s.unixClient.Post("http://_/topics/test.4",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
	}
	svc.Stop()
	svc.Wait4Stop()
	offsetsAfter := s.tkc.getOffsets("test.4")
	// Then
	delta0 := offsetsAfter[0] - offsetsBefore[0]
	delta1 := offsetsAfter[1] - offsetsBefore[1]
	imbalance := int(math.Abs(float64(delta1 - delta0)))
	if imbalance > 20 {
		panic(fmt.Errorf("Too high imbalance: %v != %v", delta0, delta1))
	}
}

// If `key` of a produced message is empty then it is deterministically
// submitted to a particular partition determined by the empty key hash.
func (s *ServiceSuite) TestProduceEmptyKey(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	for i := 0; i < 10; i++ {
		s.unixClient.Post("http://_/topics/test.4?key=",
			"text/plain", strings.NewReader(strconv.Itoa(i)))
	}
	svc.Stop()
	svc.Wait4Stop()
	offsetsAfter := s.tkc.getOffsets("test.4")
	// Then
	c.Assert(offsetsAfter[0], Equals, offsetsBefore[0])
	c.Assert(offsetsAfter[1], Equals, offsetsBefore[1])
	c.Assert(offsetsAfter[2], Equals, offsetsBefore[2])
	c.Assert(offsetsAfter[3], Equals, offsetsBefore[3]+10)
}

// Utf8 messages are submitted without a problem.
func (s *ServiceSuite) TestUtf8Message(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	s.unixClient.Post("http://_/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Превед Медвед"))
	svc.Stop()
	svc.Wait4Stop()
	// Then
	offsetsAfter := s.tkc.getOffsets("test.4")
	msgs := s.tkc.getMessages("test.4", offsetsBefore, offsetsAfter)
	c.Assert(msgs, DeepEquals,
		[][]string{[]string(nil), []string{"Превед Медвед"}, []string(nil), []string(nil)})
}

// TCP API is not started by default.
func (s *ServiceSuite) TestTCPDoesNotWork(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	// When
	r, err := s.tcpClient.Post("http://localhost:55501/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Hello Kitty"))
	// Then
	svc.Stop()
	svc.Wait4Stop()
	c.Assert(err.Error(), Matches,
		"Post http://localhost:55501/topics/test.4\\?key=foo: .* connection refused")
	c.Assert(r, IsNil)
}

// API is served on a TCP socket if it is explicitly configured.
func (s *ServiceSuite) TestBothAPI(c *C) {
	offsetsBefore := s.tkc.getOffsets("test.4")
	s.serviceCfg.TCPAddr = "127.0.0.1:55502"
	svc, _ := SpawnService(s.serviceCfg)
	// When
	_, err1 := s.tcpClient.Post("http://localhost:55502/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Превед"))
	_, err2 := s.unixClient.Post("http://_/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Kitty"))
	// Then
	svc.Stop()
	svc.Wait4Stop()
	c.Assert(err1, IsNil)
	c.Assert(err2, IsNil)
	offsetsAfter := s.tkc.getOffsets("test.4")
	msgs := s.tkc.getMessages("test.4", offsetsBefore, offsetsAfter)
	c.Assert(msgs, DeepEquals,
		[][]string{[]string(nil), []string{"Превед", "Kitty"}, []string(nil), []string(nil)})
}

func (s *ServiceSuite) TestStoppedServerCall(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	_, err := s.unixClient.Post("http://_/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Hello"))
	c.Assert(err, IsNil)
	// When
	svc.Stop()
	svc.Wait4Stop()
	// Then
	r, err := s.unixClient.Post("http://_/topics/test.4?key=foo",
		"text/plain", strings.NewReader("Kitty"))
	c.Assert(err.Error(), Equals, "Post http://_/topics/test.4?key=foo: EOF")
	c.Assert(r, IsNil)
}

// If the TCP API Server crashes then the service terminates gracefully.
func (s *ServiceSuite) TestTCPServerCrash(c *C) {
	s.serviceCfg.TCPAddr = "127.0.0.1:55502"
	svc, _ := SpawnService(s.serviceCfg)
	// When
	svc.tcpServer.errorCh <- fmt.Errorf("Kaboom!")
	// Then
	svc.Stop()
	svc.Wait4Stop()
}

// Messages that have maximum possible size indeed go through. Note that we
// assume that the broker's limit is the same as the producer's one or higher.
func (s *ServiceSuite) TestLargestMessage(c *C) {
	offsetsBefore := s.tkc.getOffsets("test.4")
	maxMsgSize := sarama.NewConfig().Producer.MaxMessageBytes - ProdMsgMetadataSize([]byte("foo"))
	msg := GenMessage(maxMsgSize)
	s.serviceCfg.TCPAddr = "127.0.0.1:55503"
	svc, _ := SpawnService(s.serviceCfg)
	// When
	r := PostChunked(s.tcpClient, "http://127.0.0.1:55503/topics/test.4?key=foo", msg)
	svc.Stop()
	svc.Wait4Stop()
	// Then
	AssertHTTPResp(c, r, http.StatusOK, "")
	offsetsAfter := s.tkc.getOffsets("test.4")
	messages := s.tkc.getMessages("test.4", offsetsBefore, offsetsAfter)
	readMsg := messages[1][0]
	c.Assert(readMsg, Equals, msg)
}

// Messages that are larger then producer's MaxMessageBytes size are silently
// dropped. Note that we assume that the broker's limit is the same as the
// producer's one or higher.
func (s *ServiceSuite) TestMessageTooLarge(c *C) {
	offsetsBefore := s.tkc.getOffsets("test.4")
	maxMsgSize := sarama.NewConfig().Producer.MaxMessageBytes - ProdMsgMetadataSize([]byte("foo")) + 1
	msg := GenMessage(maxMsgSize)
	s.serviceCfg.TCPAddr = "127.0.0.1:55504"
	svc, _ := SpawnService(s.serviceCfg)
	// When
	r := PostChunked(s.tcpClient, "http://127.0.0.1:55504/topics/test.4?key=foo", msg)
	svc.Stop()
	svc.Wait4Stop()
	// Then
	AssertHTTPResp(c, r, http.StatusOK, "")
	offsetsAfter := s.tkc.getOffsets("test.4")
	c.Assert(offsetsAfter, DeepEquals, offsetsBefore)
}

func (s *ServiceSuite) TestSyncProduce(c *C) {
	// Given
	svc, _ := SpawnService(s.serviceCfg)
	offsetsBefore := s.tkc.getOffsets("test.4")
	// When
	r, err := s.unixClient.Post("http://_/topics/test.4?key=1&sync",
		"text/plain", strings.NewReader("Foo"))
	svc.Stop()
	svc.Wait4Stop()
	offsetsAfter := s.tkc.getOffsets("test.4")
	// Then
	c.Assert(err, IsNil)
	AssertHTTPResp(c, r, http.StatusOK, "")
	c.Assert(offsetsAfter[0], Equals, offsetsBefore[0]+1)
}

func (s *ServiceSuite) TestSyncProduceInvalidTopic(c *C) {
	// Given
	svc, _ := SpawnService(s.serviceCfg)
	// When
	r, err := s.unixClient.Post("http://_/topics/no-such-topic?sync=true",
		"text/plain", strings.NewReader("Foo"))
	svc.Stop()
	svc.Wait4Stop()
	// Then
	c.Assert(err, IsNil)
	AssertHTTPResp(c, r, http.StatusNotFound,
		fmt.Sprintf("%s\n", sarama.ErrUnknownTopicOrPartition.Error()))
}

// If the Unix API Server crashes then the service terminates gracefully.
func (s *ServiceSuite) TestUnixServerCrash(c *C) {
	svc, _ := SpawnService(s.serviceCfg)
	// When
	svc.unixServer.errorCh <- fmt.Errorf("Kaboom!")
	// Then
	svc.Stop()
	svc.Wait4Stop()
}
<|endoftext|>


<commit_before>package image

import (
	"fmt"
	"os"

	"strings"

	"github.com/containerum/chkit/pkg/cli/prerun"
	"github.com/containerum/chkit/pkg/context"
	"github.com/containerum/chkit/pkg/util/activekit"
	"github.com/containerum/chkit/pkg/util/validation"
	"github.com/containerum/kube-client/pkg/model"
	"github.com/go-siris/siris/core/errors"
	"github.com/octago/sflags/gen/gpflag"
	"github.com/spf13/cobra"
)

var setAliases = []string{"imgs", "img", "im", "images"}

func Set(ctx *context.Context) *cobra.Command {
	var flags struct {
		Force      bool   `desc:"suppress confirmation" flag:"force f"`
		Deployment string `desc:"deployment name"`
		Container  string `desc:"container name"`
		Image      string `desc:"new image"`
	}
	var buildImage = func() (deployment string, image model.UpdateImage, err error) {
		var errs []string
		if validation.ValidateImageName(flags.Image) != nil {
			errs = append(errs, fmt.Sprintf("invalid image name %q", flags.Image))
		}
		if validation.ValidateContainerName(flags.Container) != nil {
			errs = append(errs, fmt.Sprintf("invalid container name %q", flags.Container))
		}
		if validation.ValidateLabel(flags.Deployment) != nil {
			errs = append(errs, fmt.Sprintf("invalid deployment name %q", flags.Deployment))
		}
		if len(errs) > 0 {
			return "", model.UpdateImage{}, errors.New(strings.Join(errs, "\n"))
		}
		return flags.Deployment, model.UpdateImage{
			Image:     flags.Image,
			Container: flags.Container,
		}, nil
	}
	command := &cobra.Command{
		Use:     "image",
		Aliases: setAliases,
		Short:   "set container image in specific deployment",
		Long: "Sets container image in specific deployment.\n" +
			"If deployment contains only one container, then uses that container by default.",
		PreRun: prerun.PreRunFunc(ctx),
		Run: func(cmd *cobra.Command, args []string) {
			var logger = ctx.Log.Command("set image")
			logger.Debugf("START")
			defer logger.Debugf("END")
			logger.StructFields(flags)
			if flags.Force {
				logger.Debugf("run command with forcr")
				var depl, image, err = buildImage()
				if err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				if err := ctx.Client.SetContainerImage(ctx.Namespace.ID, depl, image); err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				fmt.Println("OK")
				return
			}
			if flags.Deployment == "" {
				var deplList, err = ctx.Client.GetDeploymentList(ctx.Namespace.ID)
				if err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				(&activekit.Menu{
					Title: "Select deployment",
					Items: activekit.StringSelector(deplList.Names(), func(s string) error {
						flags.Deployment = s
						return nil
					}),
				}).Run()
			}
			if flags.Container == "" {
				var depl, err = ctx.Client.GetDeployment(ctx.Namespace.ID, flags.Deployment)
				if err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				if len(depl.Containers) == 1 {
					flags.Container = depl.Containers[0].Name
				} else {
					(&activekit.Menu{
						Title: "Select container",
						Items: activekit.StringSelector(depl.Containers.Names(), func(s string) error {
							flags.Container = s
							return nil
						}),
					}).Run()
				}
			}
			if flags.Image == "" {
				for {
					var image = activekit.Promt("Type new image: ")
					image = strings.TrimSpace(image)
					if validation.ValidateImageName(image) != nil {
						fmt.Printf("%q is invalid image name!\n", image)
						continue
					}
					flags.Image = image
					break
				}
			}
			if activekit.YesNo("Are you sure you want to update image to %q of container %q in deployment %s/%s?",
				flags.Image, flags.Container, ctx.Namespace, flags.Deployment) {
				if err := ctx.Client.SetContainerImage(ctx.Namespace.ID, flags.Deployment, model.UpdateImage{
					Image:     flags.Image,
					Container: flags.Container,
				}); err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
			}
		},
	}
	if err := gpflag.ParseTo(&flags, command.PersistentFlags()); err != nil {
		panic(err)
	}
	return command
}
<commit_msg>add single container case in force mode<commit_after>package image

import (
	"fmt"
	"os"

	"strings"

	"github.com/containerum/chkit/pkg/cli/prerun"
	"github.com/containerum/chkit/pkg/context"
	"github.com/containerum/chkit/pkg/util/activekit"
	"github.com/containerum/chkit/pkg/util/validation"
	"github.com/containerum/kube-client/pkg/model"
	"github.com/go-siris/siris/core/errors"
	"github.com/octago/sflags/gen/gpflag"
	"github.com/spf13/cobra"
)

var setAliases = []string{"imgs", "img", "im", "images"}

func Set(ctx *context.Context) *cobra.Command {
	var flags struct {
		Force      bool   `desc:"suppress confirmation" flag:"force f"`
		Deployment string `desc:"deployment name"`
		Container  string `desc:"container name"`
		Image      string `desc:"new image"`
	}
	var buildImage = func() (deployment string, image model.UpdateImage, err error) {
		var errs []string
		if validation.ValidateImageName(flags.Image) != nil {
			errs = append(errs, fmt.Sprintf("invalid image name %q", flags.Image))
		}
		if validation.ValidateContainerName(flags.Container) != nil {
			errs = append(errs, fmt.Sprintf("invalid container name %q", flags.Container))
		}
		if validation.ValidateLabel(flags.Deployment) != nil {
			errs = append(errs, fmt.Sprintf("invalid deployment name %q", flags.Deployment))
		}
		if len(errs) > 0 {
			return "", model.UpdateImage{}, errors.New(strings.Join(errs, "\n"))
		}
		return flags.Deployment, model.UpdateImage{
			Image:     flags.Image,
			Container: flags.Container,
		}, nil
	}
	command := &cobra.Command{
		Use:     "image",
		Aliases: setAliases,
		Short:   "set container image in specific deployment",
		Long: "Sets container image in specific deployment.\n" +
			"If deployment contains only one container, then uses that container by default.",
		PreRun: prerun.PreRunFunc(ctx),
		Run: func(cmd *cobra.Command, args []string) {
			var logger = ctx.Log.Command("set image")
			logger.Debugf("START")
			defer logger.Debugf("END")
			logger.StructFields(flags)
			if flags.Force {
				logger.Debugf("run command with force")

				if flags.Container == "" {
					var depl, err = ctx.Client.GetDeployment(ctx.Namespace.ID, flags.Deployment)
					if err != nil {
						fmt.Println(err)
						os.Exit(1)
					}
					if len(depl.Containers) == 1 {
						flags.Container = depl.Containers[0].Name
					}
				}

				var depl, image, err = buildImage()
				if err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				if err := ctx.Client.SetContainerImage(ctx.Namespace.ID, depl, image); err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				fmt.Println("OK")
				return
			}
			if flags.Deployment == "" {
				var deplList, err = ctx.Client.GetDeploymentList(ctx.Namespace.ID)
				if err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				(&activekit.Menu{
					Title: "Select deployment",
					Items: activekit.StringSelector(deplList.Names(), func(s string) error {
						flags.Deployment = s
						return nil
					}),
				}).Run()
			}
			if flags.Container == "" {
				var depl, err = ctx.Client.GetDeployment(ctx.Namespace.ID, flags.Deployment)
				if err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				if len(depl.Containers) == 1 {
					flags.Container = depl.Containers[0].Name
				} else {
					(&activekit.Menu{
						Title: "Select container",
						Items: activekit.StringSelector(depl.Containers.Names(), func(s string) error {
							flags.Container = s
							return nil
						}),
					}).Run()
				}
			}
			if flags.Image == "" {
				for {
					var image = activekit.Promt("Type new image: ")
					image = strings.TrimSpace(image)
					if validation.ValidateImageName(image) != nil {
						fmt.Printf("%q is invalid image name!\n", image)
						continue
					}
					flags.Image = image
					break
				}
			}
			if activekit.YesNo("Are you sure you want to update image to %q of container %q in deployment %s/%s?",
				flags.Image, flags.Container, ctx.Namespace, flags.Deployment) {
				if err := ctx.Client.SetContainerImage(ctx.Namespace.ID, flags.Deployment, model.UpdateImage{
					Image:     flags.Image,
					Container: flags.Container,
				}); err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
			}
		},
	}
	if err := gpflag.ParseTo(&flags, command.PersistentFlags()); err != nil {
		panic(err)
	}
	return command
}
<|endoftext|>


<commit_before>package common

import (
	"time"

	v1 "k8s.io/api/core/v1"
)

// Common types and constants used by the importer and controller.
// TODO: maybe the vm cloner can use these common values

const (
	// CDILabelKey provides a constant for CDI PVC labels
	CDILabelKey = "app"
	// CDILabelValue provides a constant  for CDI PVC label values
	CDILabelValue = "containerized-data-importer"
	// CDILabelSelector provides a constant to use for the selector to identify CDI objects in list
	CDILabelSelector = CDILabelKey + "=" + CDILabelValue

	// CDIComponentLabel can be added to all CDI resources
	CDIComponentLabel = "cdi.kubevirt.io"

	// PrometheusLabel provides the label to indicate prometheus metrics are available in the pods.
	PrometheusLabel = "prometheus.kubevirt.io"

	// ImporterVolumePath provides a constant for the directory where the PV is mounted.
	ImporterVolumePath = "/data"
	// DiskImageName provides a constant for our importer/datastream_ginkgo_test and to build ImporterWritePath
	DiskImageName = "disk.img"
	// ImporterWritePath provides a constant for the cmd/cdi-importer/importer.go executable
	ImporterWritePath = ImporterVolumePath + "/" + DiskImageName
	// ImporterWriteBlockPath provides a constant for the path where the PV is mounted.
	ImporterWriteBlockPath = "/dev/blockDevice"
	// ImporterPodName provides a constant to use as a prefix for Pods created by CDI (controller only)
	ImporterPodName = "importer"
	// ImporterDataDir provides a constant for the controller pkg to use as a hardcoded path to where content is transferred to/from (controller only)
	ImporterDataDir = "/data"
	// ScratchDataDir provides a constant for the controller pkg to use as a hardcoded path to where scratch space is located.
	ScratchDataDir = "/scratch"
	// ImporterS3Host provides an S3 string used by importer/dataStream.go only
	ImporterS3Host = "s3.amazonaws.com"
	// ImporterCertDir is where the configmap containg certs will be mounted
	ImporterCertDir = "/certs"
	// DefaultPullPolicy imports k8s "IfNotPresent" string for the import_controller_gingko_test and the cdi-controller executable
	DefaultPullPolicy = string(v1.PullIfNotPresent)

	// PullPolicy provides a constant to capture our env variable "PULL_POLICY" (only used by cmd/cdi-controller/controller.go)
	PullPolicy = "PULL_POLICY"
	// ImporterSource provides a constant to capture our env variable "IMPORTER_SOURCE"
	ImporterSource = "IMPORTER_SOURCE"
	// ImporterContentType provides a constant to capture our env variable "IMPORTER_CONTENTTYPE"
	ImporterContentType = "IMPORTER_CONTENTTYPE"
	// ImporterEndpoint provides a constant to capture our env variable "IMPORTER_ENDPOINT"
	ImporterEndpoint = "IMPORTER_ENDPOINT"
	// ImporterAccessKeyID provides a constant to capture our env variable "IMPORTER_ACCES_KEY_ID"
	ImporterAccessKeyID = "IMPORTER_ACCESS_KEY_ID"
	// ImporterSecretKey provides a constant to capture our env variable "IMPORTER_SECRET_KEY"
	ImporterSecretKey = "IMPORTER_SECRET_KEY"
	// ImporterImageSize provides a constant to capture our env variable "IMPORTER_IMAGE_SIZE"
	ImporterImageSize = "IMPORTER_IMAGE_SIZE"
	// ImporterCertDirVar provides a constant to capture our env variable "IMPORTER_CERT_DIR"
	ImporterCertDirVar = "IMPORTER_CERT_DIR"
	// InsecureTLSVar provides a constant to capture our env variable "INSECURE_TLS"
	InsecureTLSVar = "INSECURE_TLS"

	// CloningLabelKey provides a constant to use as a label name for pod affinity (controller pkg only)
	CloningLabelKey = "cloning"
	// CloningLabelValue provides a constant to use as a label value for pod affinity (controller pkg only)
	CloningLabelValue = "host-assisted-cloning"
	// CloningTopologyKey  (controller pkg only)
	CloningTopologyKey = "kubernetes.io/hostname"
	// ClonerSourcePodName (controller pkg only)
	ClonerSourcePodName = "clone-source-pod"
	// ClonerTargetPodName (controller pkg only)
	ClonerTargetPodName = "clone-target-pod"
	// ClonerImagePath (controller pkg only)
	ClonerImagePath = "/tmp/clone/image"
	// ClonerSocketPath (controller pkg only)
	ClonerSocketPath = "/tmp/clone/socket"

	// UploadServerCDILabel is the label applied to upload server resources
	UploadServerCDILabel = "cdi-upload-server"
	// UploadServerPodname is name of the upload server pod container
	UploadServerPodname = UploadServerCDILabel
	// UploadServerDataDir is the destination directoryfor uploads
	UploadServerDataDir = ImporterDataDir
	// UploadServerServiceLabel is the label selector for upload server services
	UploadServerServiceLabel = "service"
	// UploadImageSize provides a constant to capture our env variable "UPLOAD_IMAGE_SIZE"
	UploadImageSize = "UPLOAD_IMAGE_SIZE"

	// ConfigName is the name of default CDI Config
	ConfigName = "config"

	// OwnerUID provides the UID of the owner entity (either PVC or DV)
	OwnerUID = "OWNER_UID"

	// KeyAccess provides a constant to the accessKeyId label using in controller pkg and transport_test.go
	KeyAccess = "accessKeyId"
	// KeySecret provides a constant to the secretKey label using in controller pkg and transport_test.go
	KeySecret = "secretKey"

	// DefaultResyncPeriod sets a 10 minute resync period, used in the controller pkg and the controller cmd executable
	DefaultResyncPeriod = 10 * time.Minute
	// InsecureRegistryConfigMap is the name of the ConfigMap for insecure registries
	InsecureRegistryConfigMap = "cdi-insecure-registries"

	// ScratchSpaceNeededExitCode is the exit code that indicates the importer pod requires scratch space to function properly.
	ScratchSpaceNeededExitCode = 42
)
<commit_msg>Change the prometheus label used by cdi.<commit_after>package common

import (
	"time"

	v1 "k8s.io/api/core/v1"
)

// Common types and constants used by the importer and controller.
// TODO: maybe the vm cloner can use these common values

const (
	// CDILabelKey provides a constant for CDI PVC labels
	CDILabelKey = "app"
	// CDILabelValue provides a constant  for CDI PVC label values
	CDILabelValue = "containerized-data-importer"
	// CDILabelSelector provides a constant to use for the selector to identify CDI objects in list
	CDILabelSelector = CDILabelKey + "=" + CDILabelValue

	// CDIComponentLabel can be added to all CDI resources
	CDIComponentLabel = "cdi.kubevirt.io"

	// PrometheusLabel provides the label to indicate prometheus metrics are available in the pods.
	PrometheusLabel = "prometheus.cdi.kubevirt.io"

	// ImporterVolumePath provides a constant for the directory where the PV is mounted.
	ImporterVolumePath = "/data"
	// DiskImageName provides a constant for our importer/datastream_ginkgo_test and to build ImporterWritePath
	DiskImageName = "disk.img"
	// ImporterWritePath provides a constant for the cmd/cdi-importer/importer.go executable
	ImporterWritePath = ImporterVolumePath + "/" + DiskImageName
	// ImporterWriteBlockPath provides a constant for the path where the PV is mounted.
	ImporterWriteBlockPath = "/dev/blockDevice"
	// ImporterPodName provides a constant to use as a prefix for Pods created by CDI (controller only)
	ImporterPodName = "importer"
	// ImporterDataDir provides a constant for the controller pkg to use as a hardcoded path to where content is transferred to/from (controller only)
	ImporterDataDir = "/data"
	// ScratchDataDir provides a constant for the controller pkg to use as a hardcoded path to where scratch space is located.
	ScratchDataDir = "/scratch"
	// ImporterS3Host provides an S3 string used by importer/dataStream.go only
	ImporterS3Host = "s3.amazonaws.com"
	// ImporterCertDir is where the configmap containg certs will be mounted
	ImporterCertDir = "/certs"
	// DefaultPullPolicy imports k8s "IfNotPresent" string for the import_controller_gingko_test and the cdi-controller executable
	DefaultPullPolicy = string(v1.PullIfNotPresent)

	// PullPolicy provides a constant to capture our env variable "PULL_POLICY" (only used by cmd/cdi-controller/controller.go)
	PullPolicy = "PULL_POLICY"
	// ImporterSource provides a constant to capture our env variable "IMPORTER_SOURCE"
	ImporterSource = "IMPORTER_SOURCE"
	// ImporterContentType provides a constant to capture our env variable "IMPORTER_CONTENTTYPE"
	ImporterContentType = "IMPORTER_CONTENTTYPE"
	// ImporterEndpoint provides a constant to capture our env variable "IMPORTER_ENDPOINT"
	ImporterEndpoint = "IMPORTER_ENDPOINT"
	// ImporterAccessKeyID provides a constant to capture our env variable "IMPORTER_ACCES_KEY_ID"
	ImporterAccessKeyID = "IMPORTER_ACCESS_KEY_ID"
	// ImporterSecretKey provides a constant to capture our env variable "IMPORTER_SECRET_KEY"
	ImporterSecretKey = "IMPORTER_SECRET_KEY"
	// ImporterImageSize provides a constant to capture our env variable "IMPORTER_IMAGE_SIZE"
	ImporterImageSize = "IMPORTER_IMAGE_SIZE"
	// ImporterCertDirVar provides a constant to capture our env variable "IMPORTER_CERT_DIR"
	ImporterCertDirVar = "IMPORTER_CERT_DIR"
	// InsecureTLSVar provides a constant to capture our env variable "INSECURE_TLS"
	InsecureTLSVar = "INSECURE_TLS"

	// CloningLabelKey provides a constant to use as a label name for pod affinity (controller pkg only)
	CloningLabelKey = "cloning"
	// CloningLabelValue provides a constant to use as a label value for pod affinity (controller pkg only)
	CloningLabelValue = "host-assisted-cloning"
	// CloningTopologyKey  (controller pkg only)
	CloningTopologyKey = "kubernetes.io/hostname"
	// ClonerSourcePodName (controller pkg only)
	ClonerSourcePodName = "clone-source-pod"
	// ClonerTargetPodName (controller pkg only)
	ClonerTargetPodName = "clone-target-pod"
	// ClonerImagePath (controller pkg only)
	ClonerImagePath = "/tmp/clone/image"
	// ClonerSocketPath (controller pkg only)
	ClonerSocketPath = "/tmp/clone/socket"

	// UploadServerCDILabel is the label applied to upload server resources
	UploadServerCDILabel = "cdi-upload-server"
	// UploadServerPodname is name of the upload server pod container
	UploadServerPodname = UploadServerCDILabel
	// UploadServerDataDir is the destination directoryfor uploads
	UploadServerDataDir = ImporterDataDir
	// UploadServerServiceLabel is the label selector for upload server services
	UploadServerServiceLabel = "service"
	// UploadImageSize provides a constant to capture our env variable "UPLOAD_IMAGE_SIZE"
	UploadImageSize = "UPLOAD_IMAGE_SIZE"

	// ConfigName is the name of default CDI Config
	ConfigName = "config"

	// OwnerUID provides the UID of the owner entity (either PVC or DV)
	OwnerUID = "OWNER_UID"

	// KeyAccess provides a constant to the accessKeyId label using in controller pkg and transport_test.go
	KeyAccess = "accessKeyId"
	// KeySecret provides a constant to the secretKey label using in controller pkg and transport_test.go
	KeySecret = "secretKey"

	// DefaultResyncPeriod sets a 10 minute resync period, used in the controller pkg and the controller cmd executable
	DefaultResyncPeriod = 10 * time.Minute
	// InsecureRegistryConfigMap is the name of the ConfigMap for insecure registries
	InsecureRegistryConfigMap = "cdi-insecure-registries"

	// ScratchSpaceNeededExitCode is the exit code that indicates the importer pod requires scratch space to function properly.
	ScratchSpaceNeededExitCode = 42
)
<|endoftext|>


<commit_before>package couchdb

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/cozy/echo"
)

// Proxy generate a httputil.ReverseProxy which forwards the request to the
// correct route.
func Proxy(db Database, doctype, path string) *httputil.ReverseProxy {
	couchURL := config.CouchURL()
	couchAuth := config.GetConfig().CouchDB.Auth

	director := func(req *http.Request) {
		req.URL.Scheme = couchURL.Scheme
		req.URL.Host = couchURL.Host
		req.Header.Del(echo.HeaderAuthorization) // drop stack auth
		req.Header.Del(echo.HeaderCookie)
		req.URL.RawPath = "/" + makeDBName(db, doctype) + "/" + path
		req.URL.Path, _ = url.PathUnescape(req.URL.RawPath)
		if couchAuth != nil {
			if p, ok := couchAuth.Password(); ok {
				req.SetBasicAuth(couchAuth.Username(), p)
			}
		}
	}

	return &httputil.ReverseProxy{
		Director: director,
	}
}

// ProxyBulkDocs generates a httputil.ReverseProxy to forward the couchdb
// request on the _bulk_docs endpoint. This endpoint is specific since it will
// mutate many document in database, the stack has to read the response from
// couch to emit the correct realtime events.
func ProxyBulkDocs(db Database, doctype string, req *http.Request) (*httputil.ReverseProxy, *http.Request, error) {
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return nil, nil, err
	}

	var reqValue struct {
		Docs     []JSONDoc `json:"docs"`
		NewEdits *bool     `json:"new_edits"`
	}

	if err = json.Unmarshal(body, &reqValue); err != nil {
		return nil, nil, echo.NewHTTPError(http.StatusBadRequest,
			"request body is not valid JSON")
	}

	docs := make(map[string]JSONDoc)
	for _, d := range reqValue.Docs {
		docs[d.ID()] = d
	}

	// reset body to proxy
	req.Body = ioutil.NopCloser(bytes.NewReader(body))

	var transport http.RoundTripper
	if client := config.GetConfig().CouchDB.Client; client != nil {
		transport = client.Transport
	}

	p := Proxy(db, doctype, "/_bulk_docs")
	p.Transport = &bulkTransport{
		RoundTripper: transport,
		OnResponseRead: func(data []byte) {
			type respValue struct {
				ID    string `json:"id"`
				Rev   string `json:"rev"`
				OK    bool   `json:"ok"`
				Error string `json:"error"`
			}

			// When using the 'new_edits' flag (like pouchdb), the couchdb response
			// does not contain any value. We only rely on the request data and
			// expect no error.
			if reqValue.NewEdits != nil && !*reqValue.NewEdits {
				for _, doc := range reqValue.Docs {
					rev := doc.Rev()
					var event string
					if strings.HasPrefix(rev, "1-") {
						event = realtime.EventCreate
					} else {
						event = realtime.EventUpdate
					}
					RTEvent(db, event, doc, nil)
				}
			} else {
				var respValues []*respValue
				if err = json.Unmarshal(data, &respValues); err != nil {
					return
				}

				docs := make(map[string]JSONDoc)
				for _, doc := range reqValue.Docs {
					docs[doc.ID()] = doc
				}

				for _, r := range respValues {
					if r.Error != "" || !r.OK {
						continue
					}
					doc, ok := docs[r.ID]
					if !ok {
						continue
					}
					var event string
					if doc.Rev() == "" {
						event = realtime.EventCreate
					} else if doc.Get("_deleted") == true {
						event = realtime.EventDelete
					} else {
						event = realtime.EventUpdate
					}
					doc.SetRev(r.Rev)
					RTEvent(db, event, doc, nil)
				}
			}
		},
	}

	return p, req, nil
}

type bulkTransport struct {
	http.RoundTripper
	OnResponseRead func([]byte)
}

func (t *bulkTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	resp, err = t.RoundTripper.RoundTrip(req)
	if err != nil {
		return nil, newConnectionError(err)
	}
	defer func() {
		if errc := resp.Body.Close(); err == nil && errc != nil {
			err = errc
		}
	}()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode == http.StatusCreated {
		go t.OnResponseRead(b)
	}
	resp.Body = ioutil.NopCloser(bytes.NewReader(b))
	return resp, nil
}
<commit_msg>Fix transport for couch Proxy<commit_after>package couchdb

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/cozy/echo"
)

// Proxy generate a httputil.ReverseProxy which forwards the request to the
// correct route.
func Proxy(db Database, doctype, path string) *httputil.ReverseProxy {
	couchURL := config.CouchURL()
	couchAuth := config.GetConfig().CouchDB.Auth

	director := func(req *http.Request) {
		req.URL.Scheme = couchURL.Scheme
		req.URL.Host = couchURL.Host
		req.Header.Del(echo.HeaderAuthorization) // drop stack auth
		req.Header.Del(echo.HeaderCookie)
		req.URL.RawPath = "/" + makeDBName(db, doctype) + "/" + path
		req.URL.Path, _ = url.PathUnescape(req.URL.RawPath)
		if couchAuth != nil {
			if p, ok := couchAuth.Password(); ok {
				req.SetBasicAuth(couchAuth.Username(), p)
			}
		}
	}

	var transport http.RoundTripper
	if client := config.GetConfig().CouchDB.Client; client != nil {
		transport = client.Transport
	} else {
		transport = http.DefaultTransport
	}

	return &httputil.ReverseProxy{
		Director:  director,
		Transport: transport,
	}
}

// ProxyBulkDocs generates a httputil.ReverseProxy to forward the couchdb
// request on the _bulk_docs endpoint. This endpoint is specific since it will
// mutate many document in database, the stack has to read the response from
// couch to emit the correct realtime events.
func ProxyBulkDocs(db Database, doctype string, req *http.Request) (*httputil.ReverseProxy, *http.Request, error) {
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return nil, nil, err
	}

	var reqValue struct {
		Docs     []JSONDoc `json:"docs"`
		NewEdits *bool     `json:"new_edits"`
	}

	if err = json.Unmarshal(body, &reqValue); err != nil {
		return nil, nil, echo.NewHTTPError(http.StatusBadRequest,
			"request body is not valid JSON")
	}

	docs := make(map[string]JSONDoc)
	for _, d := range reqValue.Docs {
		docs[d.ID()] = d
	}

	// reset body to proxy
	req.Body = ioutil.NopCloser(bytes.NewReader(body))

	var transport http.RoundTripper
	if client := config.GetConfig().CouchDB.Client; client != nil {
		transport = client.Transport
	}

	p := Proxy(db, doctype, "/_bulk_docs")
	p.Transport = &bulkTransport{
		RoundTripper: transport,
		OnResponseRead: func(data []byte) {
			type respValue struct {
				ID    string `json:"id"`
				Rev   string `json:"rev"`
				OK    bool   `json:"ok"`
				Error string `json:"error"`
			}

			// When using the 'new_edits' flag (like pouchdb), the couchdb response
			// does not contain any value. We only rely on the request data and
			// expect no error.
			if reqValue.NewEdits != nil && !*reqValue.NewEdits {
				for _, doc := range reqValue.Docs {
					rev := doc.Rev()
					var event string
					if strings.HasPrefix(rev, "1-") {
						event = realtime.EventCreate
					} else {
						event = realtime.EventUpdate
					}
					RTEvent(db, event, doc, nil)
				}
			} else {
				var respValues []*respValue
				if err = json.Unmarshal(data, &respValues); err != nil {
					return
				}

				docs := make(map[string]JSONDoc)
				for _, doc := range reqValue.Docs {
					docs[doc.ID()] = doc
				}

				for _, r := range respValues {
					if r.Error != "" || !r.OK {
						continue
					}
					doc, ok := docs[r.ID]
					if !ok {
						continue
					}
					var event string
					if doc.Rev() == "" {
						event = realtime.EventCreate
					} else if doc.Get("_deleted") == true {
						event = realtime.EventDelete
					} else {
						event = realtime.EventUpdate
					}
					doc.SetRev(r.Rev)
					RTEvent(db, event, doc, nil)
				}
			}
		},
	}

	return p, req, nil
}

type bulkTransport struct {
	http.RoundTripper
	OnResponseRead func([]byte)
}

func (t *bulkTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	resp, err = t.RoundTripper.RoundTrip(req)
	if err != nil {
		return nil, newConnectionError(err)
	}
	defer func() {
		if errc := resp.Body.Close(); err == nil && errc != nil {
			err = errc
		}
	}()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode == http.StatusCreated {
		go t.OnResponseRead(b)
	}
	resp.Body = ioutil.NopCloser(bytes.NewReader(b))
	return resp, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package csource

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"

	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/prog"
	"github.com/google/syzkaller/sys/targets"
)

// Build builds a C/C++ program from source src and returns name of the resulting binary.
// lang can be "c" or "c++".
func Build(target *prog.Target, lang, src string) (string, error) {
	bin, err := ioutil.TempFile("", "syzkaller")
	if err != nil {
		return "", fmt.Errorf("failed to create temp file: %v", err)
	}
	bin.Close()
	sysTarget := targets.List[target.OS][target.Arch]
	compiler := sysTarget.CCompilerPrefix + "gcc"
	if _, err := exec.LookPath(compiler); err != nil {
		return "", ErrNoCompiler
	}
	flags := []string{
		"-x", lang, "-Wall", "-Werror", "-O1", "-g", "-o", bin.Name(),
		src, "-pthread",
	}
	flags = append(flags, sysTarget.CrossCFlags...)
	if sysTarget.PtrSize == 4 {
		// We do generate uint64's for syscall arguments that overflow longs on 32-bit archs.
		flags = append(flags, "-Wno-overflow")
	}
	out, err := osutil.Command(compiler, append(flags, "-static")...).CombinedOutput()
	if err != nil {
		// Some distributions don't have static libraries.
		out, err = osutil.Command(compiler, flags...).CombinedOutput()
	}
	if err != nil {
		os.Remove(bin.Name())
		data, _ := ioutil.ReadFile(src)
		return "", fmt.Errorf("failed to build program:\n%s\n%s\ncompiler invocation: %v %v",
			data, out, compiler, flags)
	}
	return bin.Name(), nil
}

var ErrNoCompiler = errors.New("no target compiler")

// Format reformats C source using clang-format.
func Format(src []byte) ([]byte, error) {
	stdout, stderr := new(bytes.Buffer), new(bytes.Buffer)
	cmd := osutil.Command("clang-format", "-assume-filename=/src.c", "-style", style)
	cmd.Stdin = bytes.NewReader(src)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	if err := cmd.Run(); err != nil {
		return src, fmt.Errorf("failed to format source: %v\n%v", err, stderr.String())
	}
	return stdout.Bytes(), nil
}

// Something acceptable for kernel developers and email-friendly.
var style = `{
BasedOnStyle: LLVM,
IndentWidth: 2,
UseTab: Never,
BreakBeforeBraces: Linux,
IndentCaseLabels: false,
DerivePointerAlignment: false,
PointerAlignment: Left,
AlignTrailingComments: true,
AllowShortBlocksOnASingleLine: false,
AllowShortCaseLabelsOnASingleLine: false,
AllowShortFunctionsOnASingleLine: false,
AllowShortIfStatementsOnASingleLine: false,
AllowShortLoopsOnASingleLine: false,
ColumnLimit: 80,
}`
<commit_msg>pkg/scource: rename compiled binary to syz-executor<commit_after>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package csource

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"

	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/prog"
	"github.com/google/syzkaller/sys/targets"
)

// Build builds a C/C++ program from source src and returns name of the resulting binary.
// lang can be "c" or "c++".
func Build(target *prog.Target, lang, src string) (string, error) {
	// We call the binary syz-executor because it sometimes shows in bug titles,
	// and we don't want 2 different bugs for when a crash is triggered during fuzzing and during repro.
	bin, err := ioutil.TempFile("", "syz-executor")
	if err != nil {
		return "", fmt.Errorf("failed to create temp file: %v", err)
	}
	bin.Close()
	sysTarget := targets.List[target.OS][target.Arch]
	compiler := sysTarget.CCompilerPrefix + "gcc"
	if _, err := exec.LookPath(compiler); err != nil {
		return "", ErrNoCompiler
	}
	flags := []string{
		"-x", lang, "-Wall", "-Werror", "-O1", "-g", "-o", bin.Name(),
		src, "-pthread",
	}
	flags = append(flags, sysTarget.CrossCFlags...)
	if sysTarget.PtrSize == 4 {
		// We do generate uint64's for syscall arguments that overflow longs on 32-bit archs.
		flags = append(flags, "-Wno-overflow")
	}
	out, err := osutil.Command(compiler, append(flags, "-static")...).CombinedOutput()
	if err != nil {
		// Some distributions don't have static libraries.
		out, err = osutil.Command(compiler, flags...).CombinedOutput()
	}
	if err != nil {
		os.Remove(bin.Name())
		data, _ := ioutil.ReadFile(src)
		return "", fmt.Errorf("failed to build program:\n%s\n%s\ncompiler invocation: %v %v",
			data, out, compiler, flags)
	}
	return bin.Name(), nil
}

var ErrNoCompiler = errors.New("no target compiler")

// Format reformats C source using clang-format.
func Format(src []byte) ([]byte, error) {
	stdout, stderr := new(bytes.Buffer), new(bytes.Buffer)
	cmd := osutil.Command("clang-format", "-assume-filename=/src.c", "-style", style)
	cmd.Stdin = bytes.NewReader(src)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	if err := cmd.Run(); err != nil {
		return src, fmt.Errorf("failed to format source: %v\n%v", err, stderr.String())
	}
	return stdout.Bytes(), nil
}

// Something acceptable for kernel developers and email-friendly.
var style = `{
BasedOnStyle: LLVM,
IndentWidth: 2,
UseTab: Never,
BreakBeforeBraces: Linux,
IndentCaseLabels: false,
DerivePointerAlignment: false,
PointerAlignment: Left,
AlignTrailingComments: true,
AllowShortBlocksOnASingleLine: false,
AllowShortCaseLabelsOnASingleLine: false,
AllowShortFunctionsOnASingleLine: false,
AllowShortIfStatementsOnASingleLine: false,
AllowShortLoopsOnASingleLine: false,
ColumnLimit: 80,
}`
<|endoftext|>


<commit_before>package domain

import (
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
)

// TODO: Add Redis as a store for the permissions

type Access interface {
	AllowedByOwner(projectID, user string) bool
	GetFile(apikey, fileID string) (*schema.File, error)
}

// access validates access to data. It checks if a user
// has been given permission to access a particular item.
type access struct {
	projects dai.Projects
	files    dai.Files
	users    dai.Users
}

// NewAccess creates a new Access.
func NewAccess(projects dai.Projects, files dai.Files, users dai.Users) *access {
	return &access{
		projects: projects,
		files:    files,
		users:    users,
	}
}

// AllowedByOwner checks to see if the user making the request has access to the
// particular item.
func (a *access) AllowedByOwner(projectID, user string) bool {
	u, err := a.users.ByID(user)
	if err != nil {
		return false
	}

	if u.Admin {
		return true
	}

	accessList, err := a.projects.AccessList(projectID)
	if err != nil {
		return false
	}
	for _, entry := range accessList {
		if user == entry.UserID {
			return true
		}
	}
	return false
}

// GetFile will validate access to a file. Rather than taking a user,
// it takes an apikey and looks up the user. It returns the file if
// access has been granted, otherwise it returns the erro ErrNoAccess.
func (a *access) GetFile(apikey, fileID string) (*schema.File, error) {
	user, err := a.users.ByAPIKey(apikey)
	if err != nil {
		// log error here
		app.Log.Error("User lookup failed", "error", err, "apikey", apikey)
		return nil, app.ErrNoAccess
	}

	file, err := a.files.ByID(fileID)
	if err != nil {
		app.Log.Error("File lookup failed", "error", err, "fileid", fileID)
		return nil, app.ErrNoAccess
	}

	if a.isInPublishedDataset(fileID) {
		return file, nil
	}

	project, err := a.files.GetProject(fileID)
	if err != nil {
		app.Log.Error("Project lookup for file failed", "error", err, "fileid", fileID)
		return nil, app.ErrNoAccess
	}

	if !a.AllowedByOwner(project.ID, user.ID) {
		app.Log.Info("Access denied", "fileid", file.ID, "user", user.ID, "projectid", project.ID)
		return nil, app.ErrNoAccess
	}

	return file, nil
}

// isInPublishedDataset checks if the file is in a published dataset. If it is then the file is accessible.
func (a *access) isInPublishedDataset(fileID string) bool {
	if datasets, err := a.files.FileDatasets(fileID); err != nil {
		return false
	} else if datasets == nil {
		return false
	} else {
		for _, ds := range datasets {
			if ds.Published {
				return true
			}
		}
		return false
	}

}
<commit_msg>Simplify isInPublishedDataset<commit_after>package domain

import (
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
)

// TODO: Add Redis as a store for the permissions

type Access interface {
	AllowedByOwner(projectID, user string) bool
	GetFile(apikey, fileID string) (*schema.File, error)
}

// access validates access to data. It checks if a user
// has been given permission to access a particular item.
type access struct {
	projects dai.Projects
	files    dai.Files
	users    dai.Users
}

// NewAccess creates a new Access.
func NewAccess(projects dai.Projects, files dai.Files, users dai.Users) *access {
	return &access{
		projects: projects,
		files:    files,
		users:    users,
	}
}

// AllowedByOwner checks to see if the user making the request has access to the
// particular item.
func (a *access) AllowedByOwner(projectID, user string) bool {
	u, err := a.users.ByID(user)
	if err != nil {
		return false
	}

	if u.Admin {
		return true
	}

	accessList, err := a.projects.AccessList(projectID)
	if err != nil {
		return false
	}
	for _, entry := range accessList {
		if user == entry.UserID {
			return true
		}
	}
	return false
}

// GetFile will validate access to a file. Rather than taking a user,
// it takes an apikey and looks up the user. It returns the file if
// access has been granted, otherwise it returns the erro ErrNoAccess.
func (a *access) GetFile(apikey, fileID string) (*schema.File, error) {
	user, err := a.users.ByAPIKey(apikey)
	if err != nil {
		// log error here
		app.Log.Error("User lookup failed", "error", err, "apikey", apikey)
		return nil, app.ErrNoAccess
	}

	file, err := a.files.ByID(fileID)
	if err != nil {
		app.Log.Error("File lookup failed", "error", err, "fileid", fileID)
		return nil, app.ErrNoAccess
	}

	if a.isInPublishedDataset(fileID) {
		return file, nil
	}

	project, err := a.files.GetProject(fileID)
	if err != nil {
		app.Log.Error("Project lookup for file failed", "error", err, "fileid", fileID)
		return nil, app.ErrNoAccess
	}

	if !a.AllowedByOwner(project.ID, user.ID) {
		app.Log.Info("Access denied", "fileid", file.ID, "user", user.ID, "projectid", project.ID)
		return nil, app.ErrNoAccess
	}

	return file, nil
}

// isInPublishedDataset checks if the file is in a published dataset. If it is then the file is accessible.
func (a *access) isInPublishedDataset(fileID string) bool {
	datasets, err := a.files.FileDatasets(fileID)
	if err != nil || datasets == nil {
		return false
	}

	for _, ds := range datasets {
		if ds.Published {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package render

import (
	"encoding/json"
	"fmt"
	"net/http"
)

// Options type to configure rendering
type Options struct {
	Status int
	Data   interface{}
	Cache  bool
}

// Renders JSON and send it to the HTTP client, supports caching
func JSON(w http.ResponseWriter, opts Options) error {
	if &w == nil {
		return fmt.Errorf("You must provide a http.ResponseWriter")
	}

	headers := w.Header()
	headers.Set("Content-Type", "application/json; charset=utf-8")

	if opts.Cache {
		headers.Set("Cache-Control", "no-cache, no-store, must-revalidate")
		headers.Set("Pragma", "no-cache")
		headers.Set("Expires", "0")
	}

	jsonbytes, err := json.Marshal(opts.Data)
	if err != nil {
		return err
	}

	headers.Set("Content-Length", fmt.Sprintf("%d", len(jsonbytes)))
	if opts.Status <= 0 {
		opts.Status = http.StatusOK
	}
	w.WriteHeader(opts.Status)

	_, err = w.Write(jsonbytes)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Cleans up code to make it look more idiomatic.<commit_after>package render

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
)

// Options type to configure rendering
type Options struct {
	Status int
	Data   interface{}
	Cache  bool
}

// Renders JSON and send it to the HTTP client, supports caching
func JSON(w http.ResponseWriter, opts Options) error {
	if &w == nil {
		return fmt.Errorf("You must provide a http.ResponseWriter")
	}

	headers := w.Header()
	headers.Set("Content-Type", "application/json; charset=utf-8")

	if opts.Cache {
		headers.Set("Cache-Control", "no-cache, no-store, must-revalidate")
		headers.Set("Pragma", "no-cache")
		headers.Set("Expires", "0")
	}

	jsonBytes, err := json.Marshal(opts.Data)
	if err != nil {
		return err
	}

	headers.Set("Content-Length", strconv.Itoa(len(jsonBytes)))
	if opts.Status <= 0 {
		opts.Status = http.StatusOK
	}
	w.WriteHeader(opts.Status)
	w.Write(jsonBytes)

	return nil
}
<|endoftext|>