<commit_before>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   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 echo

import (
	"net/http"
	"net/url"
	"time"
)

// CookieOptions cookie options
type CookieOptions struct {
	Prefix string

	// MaxAge=0 means no 'Max-Age' attribute specified.
	// MaxAge<0 means delete cookie now, equivalently 'Max-Age: 0'.
	// MaxAge>0 means Max-Age attribute present and given in seconds.
	MaxAge int

	Path     string
	Domain   string
	Secure   bool
	HttpOnly bool
}

func (c *CookieOptions) Clone() *CookieOptions {
	clone := *c
	return &clone
}

//Cookier interface
type Cookier interface {
	Get(key string) string
	Set(key string, val string, args ...interface{}) Cookier
}

//NewCookier create a cookie instance
func NewCookier(ctx Context) Cookier {
	return &cookie{
		context: ctx,
		cookies: []*Cookie{},
	}
}

//NewCookie create a cookie instance
func NewCookie(name string, value string, opts ...*CookieOptions) *Cookie {
	opt := &CookieOptions{}
	if len(opts) > 0 {
		opt = opts[0]
	}
	if len(opt.Path) == 0 {
		opt.Path = `/`
	}
	cookie := &Cookie{
		cookie: &http.Cookie{
			Name:     opt.Prefix + name,
			Value:    value,
			Path:     opt.Path,
			Domain:   opt.Domain,
			MaxAge:   opt.MaxAge,
			Secure:   opt.Secure,
			HttpOnly: opt.HttpOnly,
		},
	}
	return cookie
}

//Cookie 操作封装
type Cookie struct {
	cookie *http.Cookie
}

//Path 设置路径
func (c *Cookie) Path(p string) *Cookie {
	c.cookie.Path = p
	return c
}

//Domain 设置域名
func (c *Cookie) Domain(p string) *Cookie {
	c.cookie.Domain = p
	return c
}

//MaxAge 设置有效时长（秒）
func (c *Cookie) MaxAge(p int) *Cookie {
	c.cookie.MaxAge = p
	return c
}

//Expires 设置过期时间戳
func (c *Cookie) Expires(p int64) *Cookie {
	if p > 0 {
		c.cookie.Expires = time.Unix(time.Now().Unix()+p, 0)
	} else if p < 0 {
		c.cookie.Expires = time.Unix(1, 0)
	}
	return c
}

//Secure 设置是否启用HTTPS
func (c *Cookie) Secure(p bool) *Cookie {
	c.cookie.Secure = p
	return c
}

//HttpOnly 设置是否启用HttpOnly
func (c *Cookie) HttpOnly(p bool) *Cookie {
	c.cookie.HttpOnly = p
	return c
}

//Send 发送cookie数据到响应头
func (c *Cookie) Send(ctx Context) {
	ctx.Response().SetCookie(c.cookie)
}

type cookie struct {
	context Context
	cookies []*Cookie
}

func (c *cookie) Get(key string) string {
	var val string
	if v := c.context.Request().Cookie(c.context.CookieOptions().Prefix + key); len(v) > 0 {
		val, _ = url.QueryUnescape(v)
	}
	return val
}

func (c *cookie) Set(key string, val string, args ...interface{}) Cookier {
	val = url.QueryEscape(val)
	var cookie *Cookie
	var found bool
	for _, v := range c.cookies {
		if key == v.cookie.Name {
			cookie = v
			found = true
			break
		}
	}
	if cookie == nil {
		cookie = NewCookie(key, val, c.context.CookieOptions())
	}
	switch len(args) {
	case 5:
		httpOnly, _ := args[4].(bool)
		cookie.HttpOnly(httpOnly)
		fallthrough
	case 4:
		secure, _ := args[3].(bool)
		cookie.Secure(secure)
		fallthrough
	case 3:
		domain, _ := args[2].(string)
		cookie.Domain(domain)
		fallthrough
	case 2:
		ppath, _ := args[1].(string)
		cookie.Path(ppath)
		fallthrough
	case 1:
		var liftTime int64
		switch args[0].(type) {
		case int:
			liftTime = int64(args[0].(int))
		case int64:
			liftTime = args[0].(int64)
		case time.Duration:
			liftTime = int64(args[0].(time.Duration).Seconds())
		}
		cookie.Expires(liftTime)
	}
	if !found {
		c.cookies = append(c.cookies, cookie)
		cookie.Send(c.context)
	} else {
		c.context.Response().Header().Del(HeaderSetCookie)
		for _, cookie := range c.cookies {
			cookie.Send(c.context)
		}
	}
	return c
}
<commit_msg>cookie: 增加sameSite属性支持<commit_after>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   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 echo

import (
	"net/http"
	"net/url"
	"strings"
	"time"
)

// CookieOptions cookie options
type CookieOptions struct {
	Prefix string

	// MaxAge=0 means no 'Max-Age' attribute specified.
	// MaxAge<0 means delete cookie now, equivalently 'Max-Age: 0'.
	// MaxAge>0 means Max-Age attribute present and given in seconds.
	MaxAge int

	Path     string
	Domain   string
	Secure   bool
	HttpOnly bool
	SameSite string // strict / lax
}

func (c *CookieOptions) Clone() *CookieOptions {
	clone := *c
	return &clone
}

//Cookier interface
type Cookier interface {
	Get(key string) string
	Set(key string, val string, args ...interface{}) Cookier
}

//NewCookier create a cookie instance
func NewCookier(ctx Context) Cookier {
	return &cookie{
		context: ctx,
		cookies: []*Cookie{},
	}
}

//NewCookie create a cookie instance
func NewCookie(name string, value string, opts ...*CookieOptions) *Cookie {
	opt := &CookieOptions{}
	if len(opts) > 0 {
		opt = opts[0]
	}
	if len(opt.Path) == 0 {
		opt.Path = `/`
	}
	cookie := &Cookie{
		cookie: &http.Cookie{
			Name:     opt.Prefix + name,
			Value:    value,
			Path:     opt.Path,
			Domain:   opt.Domain,
			MaxAge:   opt.MaxAge,
			Secure:   opt.Secure,
			HttpOnly: opt.HttpOnly,
		},
	}
	if len(opt.SameSite) > 0 {
		cookie.SameSite(opt.SameSite)
	}
	return cookie
}

//Cookie 操作封装
type Cookie struct {
	cookie *http.Cookie
}

//Path 设置路径
func (c *Cookie) Path(p string) *Cookie {
	c.cookie.Path = p
	return c
}

//Domain 设置域名
func (c *Cookie) Domain(p string) *Cookie {
	c.cookie.Domain = p
	return c
}

//MaxAge 设置有效时长（秒）
func (c *Cookie) MaxAge(p int) *Cookie {
	c.cookie.MaxAge = p
	return c
}

//Expires 设置过期时间戳
func (c *Cookie) Expires(p int64) *Cookie {
	if p > 0 {
		c.cookie.Expires = time.Unix(time.Now().Unix()+p, 0)
	} else if p < 0 {
		c.cookie.Expires = time.Unix(1, 0)
	}
	return c
}

//Secure 设置是否启用HTTPS
func (c *Cookie) Secure(p bool) *Cookie {
	c.cookie.Secure = p
	return c
}

//HttpOnly 设置是否启用HttpOnly
func (c *Cookie) HttpOnly(p bool) *Cookie {
	c.cookie.HttpOnly = p
	return c
}

//SameSite 设置SameSite
func (c *Cookie) SameSite(p string) *Cookie {
	switch strings.ToLower(p) {
	case `lax`:
		c.cookie.SameSite = http.SameSiteLaxMode
	case `strict`:
		c.cookie.SameSite = http.SameSiteStrictMode
	default:
		c.cookie.SameSite = http.SameSiteDefaultMode
	}
	return c
}

//Send 发送cookie数据到响应头
func (c *Cookie) Send(ctx Context) {
	ctx.Response().SetCookie(c.cookie)
}

type cookie struct {
	context Context
	cookies []*Cookie
}

func (c *cookie) Get(key string) string {
	var val string
	if v := c.context.Request().Cookie(c.context.CookieOptions().Prefix + key); len(v) > 0 {
		val, _ = url.QueryUnescape(v)
	}
	return val
}

func (c *cookie) Set(key string, val string, args ...interface{}) Cookier {
	val = url.QueryEscape(val)
	var cookie *Cookie
	var found bool
	for _, v := range c.cookies {
		if key == v.cookie.Name {
			cookie = v
			found = true
			break
		}
	}
	if cookie == nil {
		cookie = NewCookie(key, val, c.context.CookieOptions())
	}
	switch len(args) {
	case 6:
		sameSite, _ := args[5].(string)
		cookie.SameSite(sameSite)
		fallthrough
	case 5:
		httpOnly, _ := args[4].(bool)
		cookie.HttpOnly(httpOnly)
		fallthrough
	case 4:
		secure, _ := args[3].(bool)
		cookie.Secure(secure)
		fallthrough
	case 3:
		domain, _ := args[2].(string)
		cookie.Domain(domain)
		fallthrough
	case 2:
		ppath, _ := args[1].(string)
		cookie.Path(ppath)
		fallthrough
	case 1:
		var liftTime int64
		switch args[0].(type) {
		case int:
			liftTime = int64(args[0].(int))
		case int64:
			liftTime = args[0].(int64)
		case time.Duration:
			liftTime = int64(args[0].(time.Duration).Seconds())
		}
		cookie.Expires(liftTime)
	}
	if !found {
		c.cookies = append(c.cookies, cookie)
		cookie.Send(c.context)
	} else {
		c.context.Response().Header().Del(HeaderSetCookie)
		for _, cookie := range c.cookies {
			cookie.Send(c.context)
		}
	}
	return c
}
<|endoftext|>


<commit_before>package rest

import (
	"testing"
	"github.com/maniksurtani/quotaservice/config"
	"github.com/maniksurtani/quotaservice"
	"github.com/maniksurtani/quotaservice/admin"
	"fmt"
)

func TestAddGlobalDefault(t *testing.T) {
	s := startService(false)
	defer s.Stop()

	assertDefaultBucketDoesNotExist(t, s)

	b := config.NewDefaultBucketConfig()
	e := s.(admin.Administrable).AddBucket(config.GlobalNamespace, b.ToProto())
	assertNoError(t, e)

	assertDefaultBucketExists(t, s)

	// Now try and add a bucket config again - should error.
	e = s.(admin.Administrable).AddBucket(config.GlobalNamespace, b.ToProto())
	assertError(t, e)
}

func TestRemoveGlobalDefault(t *testing.T) {
	s := startService(true)
	defer s.Stop()

	assertDefaultBucketExists(t, s)

	e := s.(admin.Administrable).DeleteBucket(config.GlobalNamespace, config.DefaultBucketName)
	assertNoError(t, e)

	assertDefaultBucketDoesNotExist(t, s)

	// Should be idempotent
	e = s.(admin.Administrable).DeleteBucket(config.GlobalNamespace, config.DefaultBucketName)
	assertNoError(t, e)

	assertDefaultBucketDoesNotExist(t, s)
}

func TestUpdateGlobalDefault(t *testing.T) {
	s := startService(true)
	defer s.Stop()

	assertDefaultBucketExists(t, s)

	b := config.NewDefaultBucketConfig()
	b.MaxTokensPerRequest = 2
	b.Name = config.DefaultBucketName
	e := s.(admin.Administrable).UpdateBucket(config.GlobalNamespace, b.ToProto())
	assertNoError(t, e)

	// Now check that we hit max tokens limits.
	_, _, e = s.(quotaservice.QuotaService).Allow("doesn't exist", "doesn't exist", 5, 0)
	assertError(t, e)
	if e.(quotaservice.QuotaServiceError).Reason != quotaservice.ER_TOO_MANY_TOKENS_REQUESTED {
		t.Fatal("Wrong error: ", e)
	}

	b.MaxTokensPerRequest = 10
	e = s.(admin.Administrable).UpdateBucket(config.GlobalNamespace, b.ToProto())
	assertNoError(t, e)

	// Now check again
	_, _, e = s.(quotaservice.QuotaService).Allow("doesn't exist", "doesn't exist", 5, 0)
	assertNoError(t, e)
}

func TestAddNamespace(t *testing.T) {
	s := startService(false)
	defer s.Stop()

	assertBucketDoesNotExist(t, s, "ns", "b")

	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	n.SetDynamicBucketTemplate(config.NewDefaultBucketConfig())

	e := s.(admin.Administrable).AddNamespace(n.ToProto())
	assertNoError(t, e)

	assertBucketExists(t, s, "ns", "b")
	assertBucketExists(t, s, "ns", "bb")
	assertBucketExists(t, s, "ns", "bbb")

	e = s.(admin.Administrable).AddNamespace(n.ToProto())
	assertError(t, e)

	assertBucketExists(t, s, "ns", "bbbb")
	assertBucketExists(t, s, "ns", "bbbbb")
	assertBucketExists(t, s, "ns", "bbbbbb")
}

func TestRemoveNamespace(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	n.SetDynamicBucketTemplate(config.NewDefaultBucketConfig())

	s := startService(false, n)
	defer s.Stop()

	assertBucketExists(t, s, "ns", "b")

	e := s.(admin.Administrable).DeleteNamespace("ns")
	assertNoError(t, e)

	assertBucketDoesNotExist(t, s, "ns", "b")

	e = s.(admin.Administrable).DeleteNamespace("ns")
	assertError(t, e)

	assertBucketDoesNotExist(t, s, "ns", "b")
}

func TestUpdateNamespace(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	n.SetDynamicBucketTemplate(config.NewDefaultBucketConfig())

	s := startService(false, n)
	defer s.Stop()

	// Allows dynamic buckets.
	assertBucketExists(t, s, "ns", "b")
	assertBucketExists(t, s, "ns", "bb")
	assertBucketExists(t, s, "ns", "bbb")

	// change config to not allow dynamic buckets
	n.DynamicBucketTemplate = nil
	e := s.(admin.Administrable).UpdateNamespace(n.ToProto())
	assertNoError(t, e)

	// Existing buckets should have been removed.
	assertBucketDoesNotExist(t, s, "ns", "b")
	assertBucketDoesNotExist(t, s, "ns", "bb")
	assertBucketDoesNotExist(t, s, "ns", "bbb")

	// No new dynamic buckets
	assertBucketDoesNotExist(t, s, "ns", "bbbb")
	assertBucketDoesNotExist(t, s, "ns", "bbbbb")
	assertBucketDoesNotExist(t, s, "ns", "bbbbbb")
}

func TestAddBucket(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	b := config.NewDefaultBucketConfig()
	b.Name = "b"
	n.AddBucket("b", b)

	s := startService(false, n)
	defer s.Stop()

	// Doesn't allow dynamic buckets.
	assertBucketExists(t, s, "ns", "b")
	assertBucketDoesNotExist(t, s, "ns", "b1")

	// Add bucket
	b.Name = "b1"
	e := s.(admin.Administrable).AddBucket("ns", b.ToProto())
	assertNoError(t, e)

	// Existing buckets should still be there
	assertBucketExists(t, s, "ns", "b")
	assertBucketExists(t, s, "ns", "b1")
	assertBucketDoesNotExist(t, s, "ns", "b2")

	// Already exists
	e = s.(admin.Administrable).AddBucket("ns", b.ToProto())
	assertError(t, e)
}

func TestRemoveBucket(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	b := config.NewDefaultBucketConfig()
	b.Name = "b"
	n.AddBucket("b", b)

	s := startService(false, n)
	defer s.Stop()

	// Doesn't allow dynamic buckets.
	assertBucketExists(t, s, "ns", "b")
	assertBucketDoesNotExist(t, s, "ns", "b1")

	// Add bucket
	e := s.(admin.Administrable).DeleteBucket("ns", "b")
	assertNoError(t, e)

	// Existing buckets should still be there
	assertBucketDoesNotExist(t, s, "ns", "b")

	// Idempotence
	e = s.(admin.Administrable).DeleteBucket("ns", "b")
	assertNoError(t, e)
}

func TestUpdateBucket(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	b := config.NewDefaultBucketConfig()
	b.Name = "b"
	b.MaxTokensPerRequest = 2
	n.AddBucket("b", b)

	s := startService(false, n)
	defer s.Stop()

	assertBucketExists(t, s, "ns", "b")

	// Now check that we hit max tokens limits.
	_, _, e := s.(quotaservice.QuotaService).Allow("ns", "b", 5, 0)
	assertError(t, e)
	if e.(quotaservice.QuotaServiceError).Reason != quotaservice.ER_TOO_MANY_TOKENS_REQUESTED {
		t.Fatal("Wrong error: ", e)
	}

	// Update bucket
	b.MaxTokensPerRequest = 10
	e = s.(admin.Administrable).UpdateBucket("ns", b.ToProto())
	assertNoError(t, e)

	_, _, e = s.(quotaservice.QuotaService).Allow("ns", "b", 5, 0)
	assertNoError(t, e)
}

func startService(withDefault bool, ns ...*config.NamespaceConfig) quotaservice.Server {
	c := config.NewDefaultServiceConfig()
	if !withDefault {
		c.GlobalDefaultBucket = nil
	}
	for _, n := range ns {
		fmt.Printf("Adding namespace %+v ", n)
		c.AddNamespace(n.Name, n)
	}
	s := quotaservice.New(c, &quotaservice.MockBucketFactory{}, &quotaservice.MockEndpoint{})
	s.Start()
	return s
}

func assertDefaultBucketExists(t *testing.T, s quotaservice.Server) {
	assertBucketExists(t, s, "doesn't exist", "doesn't exist")
}

func assertBucketExists(t *testing.T, s quotaservice.Server, nsName, bName string) {
	// Demonstrate that we now do have a default bucket.
	g, w, e := s.(quotaservice.QuotaService).Allow(nsName, bName, 1, 0)
	assertNoError(t, e)

	if g != 1 {
		t.Fatal("Expected to be granted 1 token")
	}

	if w != 0 {
		t.Fatal("Expecting wait time of 0")
	}
}

func assertDefaultBucketDoesNotExist(t *testing.T, s quotaservice.Server) {
	assertBucketDoesNotExist(t, s, "doesn't exist", "doesn't exist")
}

func assertBucketDoesNotExist(t *testing.T, s quotaservice.Server, nsName, bName string) {
	// Demonstrate that there is no default bucket first
	_, _, e := s.(quotaservice.QuotaService).Allow(nsName, bName, 1, 0)
	assertError(t, e)
}

func assertNoError(t *testing.T, e error) {
	if e != nil {
		t.Fatal("Not expecting error ", e)
	}
}

func assertError(t *testing.T, e error) {
	if e == nil {
		t.Fatal("Expecting error!")
	}
}
<commit_msg>Gofmt<commit_after>package rest

import (
	"fmt"
	"github.com/maniksurtani/quotaservice"
	"github.com/maniksurtani/quotaservice/admin"
	"github.com/maniksurtani/quotaservice/config"
	"testing"
)

func TestAddGlobalDefault(t *testing.T) {
	s := startService(false)
	defer s.Stop()

	assertDefaultBucketDoesNotExist(t, s)

	b := config.NewDefaultBucketConfig()
	e := s.(admin.Administrable).AddBucket(config.GlobalNamespace, b.ToProto())
	assertNoError(t, e)

	assertDefaultBucketExists(t, s)

	// Now try and add a bucket config again - should error.
	e = s.(admin.Administrable).AddBucket(config.GlobalNamespace, b.ToProto())
	assertError(t, e)
}

func TestRemoveGlobalDefault(t *testing.T) {
	s := startService(true)
	defer s.Stop()

	assertDefaultBucketExists(t, s)

	e := s.(admin.Administrable).DeleteBucket(config.GlobalNamespace, config.DefaultBucketName)
	assertNoError(t, e)

	assertDefaultBucketDoesNotExist(t, s)

	// Should be idempotent
	e = s.(admin.Administrable).DeleteBucket(config.GlobalNamespace, config.DefaultBucketName)
	assertNoError(t, e)

	assertDefaultBucketDoesNotExist(t, s)
}

func TestUpdateGlobalDefault(t *testing.T) {
	s := startService(true)
	defer s.Stop()

	assertDefaultBucketExists(t, s)

	b := config.NewDefaultBucketConfig()
	b.MaxTokensPerRequest = 2
	b.Name = config.DefaultBucketName
	e := s.(admin.Administrable).UpdateBucket(config.GlobalNamespace, b.ToProto())
	assertNoError(t, e)

	// Now check that we hit max tokens limits.
	_, _, e = s.(quotaservice.QuotaService).Allow("doesn't exist", "doesn't exist", 5, 0)
	assertError(t, e)
	if e.(quotaservice.QuotaServiceError).Reason != quotaservice.ER_TOO_MANY_TOKENS_REQUESTED {
		t.Fatal("Wrong error: ", e)
	}

	b.MaxTokensPerRequest = 10
	e = s.(admin.Administrable).UpdateBucket(config.GlobalNamespace, b.ToProto())
	assertNoError(t, e)

	// Now check again
	_, _, e = s.(quotaservice.QuotaService).Allow("doesn't exist", "doesn't exist", 5, 0)
	assertNoError(t, e)
}

func TestAddNamespace(t *testing.T) {
	s := startService(false)
	defer s.Stop()

	assertBucketDoesNotExist(t, s, "ns", "b")

	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	n.SetDynamicBucketTemplate(config.NewDefaultBucketConfig())

	e := s.(admin.Administrable).AddNamespace(n.ToProto())
	assertNoError(t, e)

	assertBucketExists(t, s, "ns", "b")
	assertBucketExists(t, s, "ns", "bb")
	assertBucketExists(t, s, "ns", "bbb")

	e = s.(admin.Administrable).AddNamespace(n.ToProto())
	assertError(t, e)

	assertBucketExists(t, s, "ns", "bbbb")
	assertBucketExists(t, s, "ns", "bbbbb")
	assertBucketExists(t, s, "ns", "bbbbbb")
}

func TestRemoveNamespace(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	n.SetDynamicBucketTemplate(config.NewDefaultBucketConfig())

	s := startService(false, n)
	defer s.Stop()

	assertBucketExists(t, s, "ns", "b")

	e := s.(admin.Administrable).DeleteNamespace("ns")
	assertNoError(t, e)

	assertBucketDoesNotExist(t, s, "ns", "b")

	e = s.(admin.Administrable).DeleteNamespace("ns")
	assertError(t, e)

	assertBucketDoesNotExist(t, s, "ns", "b")
}

func TestUpdateNamespace(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	n.SetDynamicBucketTemplate(config.NewDefaultBucketConfig())

	s := startService(false, n)
	defer s.Stop()

	// Allows dynamic buckets.
	assertBucketExists(t, s, "ns", "b")
	assertBucketExists(t, s, "ns", "bb")
	assertBucketExists(t, s, "ns", "bbb")

	// change config to not allow dynamic buckets
	n.DynamicBucketTemplate = nil
	e := s.(admin.Administrable).UpdateNamespace(n.ToProto())
	assertNoError(t, e)

	// Existing buckets should have been removed.
	assertBucketDoesNotExist(t, s, "ns", "b")
	assertBucketDoesNotExist(t, s, "ns", "bb")
	assertBucketDoesNotExist(t, s, "ns", "bbb")

	// No new dynamic buckets
	assertBucketDoesNotExist(t, s, "ns", "bbbb")
	assertBucketDoesNotExist(t, s, "ns", "bbbbb")
	assertBucketDoesNotExist(t, s, "ns", "bbbbbb")
}

func TestAddBucket(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	b := config.NewDefaultBucketConfig()
	b.Name = "b"
	n.AddBucket("b", b)

	s := startService(false, n)
	defer s.Stop()

	// Doesn't allow dynamic buckets.
	assertBucketExists(t, s, "ns", "b")
	assertBucketDoesNotExist(t, s, "ns", "b1")

	// Add bucket
	b.Name = "b1"
	e := s.(admin.Administrable).AddBucket("ns", b.ToProto())
	assertNoError(t, e)

	// Existing buckets should still be there
	assertBucketExists(t, s, "ns", "b")
	assertBucketExists(t, s, "ns", "b1")
	assertBucketDoesNotExist(t, s, "ns", "b2")

	// Already exists
	e = s.(admin.Administrable).AddBucket("ns", b.ToProto())
	assertError(t, e)
}

func TestRemoveBucket(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	b := config.NewDefaultBucketConfig()
	b.Name = "b"
	n.AddBucket("b", b)

	s := startService(false, n)
	defer s.Stop()

	// Doesn't allow dynamic buckets.
	assertBucketExists(t, s, "ns", "b")
	assertBucketDoesNotExist(t, s, "ns", "b1")

	// Add bucket
	e := s.(admin.Administrable).DeleteBucket("ns", "b")
	assertNoError(t, e)

	// Existing buckets should still be there
	assertBucketDoesNotExist(t, s, "ns", "b")

	// Idempotence
	e = s.(admin.Administrable).DeleteBucket("ns", "b")
	assertNoError(t, e)
}

func TestUpdateBucket(t *testing.T) {
	n := config.NewDefaultNamespaceConfig()
	n.Name = "ns"
	b := config.NewDefaultBucketConfig()
	b.Name = "b"
	b.MaxTokensPerRequest = 2
	n.AddBucket("b", b)

	s := startService(false, n)
	defer s.Stop()

	assertBucketExists(t, s, "ns", "b")

	// Now check that we hit max tokens limits.
	_, _, e := s.(quotaservice.QuotaService).Allow("ns", "b", 5, 0)
	assertError(t, e)
	if e.(quotaservice.QuotaServiceError).Reason != quotaservice.ER_TOO_MANY_TOKENS_REQUESTED {
		t.Fatal("Wrong error: ", e)
	}

	// Update bucket
	b.MaxTokensPerRequest = 10
	e = s.(admin.Administrable).UpdateBucket("ns", b.ToProto())
	assertNoError(t, e)

	_, _, e = s.(quotaservice.QuotaService).Allow("ns", "b", 5, 0)
	assertNoError(t, e)
}

func startService(withDefault bool, ns ...*config.NamespaceConfig) quotaservice.Server {
	c := config.NewDefaultServiceConfig()
	if !withDefault {
		c.GlobalDefaultBucket = nil
	}
	for _, n := range ns {
		fmt.Printf("Adding namespace %+v ", n)
		c.AddNamespace(n.Name, n)
	}
	s := quotaservice.New(c, &quotaservice.MockBucketFactory{}, &quotaservice.MockEndpoint{})
	s.Start()
	return s
}

func assertDefaultBucketExists(t *testing.T, s quotaservice.Server) {
	assertBucketExists(t, s, "doesn't exist", "doesn't exist")
}

func assertBucketExists(t *testing.T, s quotaservice.Server, nsName, bName string) {
	// Demonstrate that we now do have a default bucket.
	g, w, e := s.(quotaservice.QuotaService).Allow(nsName, bName, 1, 0)
	assertNoError(t, e)

	if g != 1 {
		t.Fatal("Expected to be granted 1 token")
	}

	if w != 0 {
		t.Fatal("Expecting wait time of 0")
	}
}

func assertDefaultBucketDoesNotExist(t *testing.T, s quotaservice.Server) {
	assertBucketDoesNotExist(t, s, "doesn't exist", "doesn't exist")
}

func assertBucketDoesNotExist(t *testing.T, s quotaservice.Server, nsName, bName string) {
	// Demonstrate that there is no default bucket first
	_, _, e := s.(quotaservice.QuotaService).Allow(nsName, bName, 1, 0)
	assertError(t, e)
}

func assertNoError(t *testing.T, e error) {
	if e != nil {
		t.Fatal("Not expecting error ", e)
	}
}

func assertError(t *testing.T, e error) {
	if e == nil {
		t.Fatal("Expecting error!")
	}
}
<|endoftext|>


<commit_before>// Package apply provides utilities to apply a project config to GCP by deploying all defined resources.
package apply

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/healthcare/deploy/config"
	"github.com/GoogleCloudPlatform/healthcare/deploy/deploymentmanager"
)

const (
	deploymentNamePrefix            = "data-protect-toolkit"
	auditDeploymentName             = deploymentNamePrefix + "-audit"
	resourceDeploymentName          = deploymentNamePrefix + "-resources"
	setupPrerequisiteDeploymentName = deploymentNamePrefix + "-prerequisites"
)

// deploymentManagerRoles are the roles granted to the DM service account.
var deploymentManagerRoles = []string{"roles/owner", "roles/storage.admin"}

// deploymentRetryWaitTime is the time to wait between retrying a deployment to allow for concurrent operations to finish.
const deploymentRetryWaitTime = time.Minute

// The following vars are stubbed in tests.
var (
	cmdOutput = func(cmd *exec.Cmd) ([]byte, error) {
		log.Printf("Running: %v", cmd.Args)
    cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return cmd.Output()
	}
	cmdRun = func(cmd *exec.Cmd) error {
		log.Printf("Running: %v", cmd.Args)
    cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return cmd.Run()
	}
	upsertDeployment         = deploymentmanager.Upsert
	upsertDeploymentFromFile = deploymentmanager.UpsertFromFile
)

// parsedResource is an interface that must be implemented by all concrete resource implementations.
type parsedResource interface {
	Init(*config.Project) error
	Name() string
}

// deploymentManagerTyper should be implemented by resources that are natively supported by the deployment manager service.
// Use this if there is no suitable CFT template for a resource and a custom template is not needed.
// See https://cloud.google.com/deployment-manager/docs/configuration/supported-resource-types for valid types.
type deploymentManagerTyper interface {
	DeploymentManagerType() string
}

// deploymentManagerPather should be implemented by resources that use a DM template to deploy.
// Use this if the resource wraps a CFT or custom template.
type deploymentManagerPather interface {
	TemplatePath() string
}

// depender is the interface that defines a method to get dependent resources.
type depender interface {
	// Dependencies returns the name of the resource IDs to depend on.
	Dependencies() []string
}

// Apply deploys the CFT resources in the project.
func Apply(conf *config.Config, project *config.Project) error {
	if err := grantDeploymentManagerAccess(project); err != nil {
		return fmt.Errorf("failed to grant deployment manager access to the project: %v", err)
	}

	if err := deployPrerequisite(project); err != nil {
		return fmt.Errorf("failed to deploy CHC Type Provider: %v", err)
	}

	// TODO: stop retrying once
	// https://github.com/GoogleCloudPlatform/cloud-foundation-toolkit/issues/17
	// is fixed.
	for i := 0; i < 2; i++ {
		if err := deployResources(project); err == nil {
			break
		} else if i == 1 {
			return fmt.Errorf("failed to deploy resources: %v", err)
		}
		log.Printf("Sleeping for %v and retrying in case failure was due to concurrent IAM policy update", deploymentRetryWaitTime)
		time.Sleep(deploymentRetryWaitTime)
	}

	// Always get the latest log sink writer as when the sink is moved between deployments it may
	// create a new sink writer.
	sinkSA, err := getLogSinkServiceAccount(project)
	if err != nil {
		return fmt.Errorf("failed to get log sink service account: %v", err)
	}

	// Note: if the project was deployed the project Init will already have added the log sink service account permission on the dataset.
	if currSA := project.GeneratedFields.LogSinkServiceAccount; currSA == "" {
		project.AuditLogs.LogsBQDataset.Accesses = append(project.AuditLogs.LogsBQDataset.Accesses, &config.Access{
			Role: "WRITER", UserByEmail: sinkSA,
		})
	} else if currSA != sinkSA {
		project.GeneratedFields.LogSinkServiceAccount = sinkSA
		// Replace all instances of old writer SA with new.
		for _, a := range project.AuditLogs.LogsBQDataset.Accesses {
			if a.UserByEmail == currSA {
				a.UserByEmail = sinkSA
			}
		}
	}

	if err := deployAudit(project, conf.ProjectForAuditLogs(project)); err != nil {
		return fmt.Errorf("failed to deploy audit resources: %v", err)
	}

	if err := deployGKEWorkloads(project); err != nil {
		return fmt.Errorf("failed to deploy GKE workloads: %v", err)
	}

	// Only remove owner account if there is an organization to ensure the project has an administrator.
	if conf.Overall.OrganizationID != "" {
		if err := removeOwnerUser(project); err != nil {
			return fmt.Errorf("failed to remove owner user: %v", err)
		}
	}
	return nil
}

// grantDeploymentManagerAccess grants the necessary permissions to the DM service account to perform its actions.
// Note: we don't revoke deployment manager's access because permissions can take up to 7 minutes
// to propagate through the system, which can cause permission denied issues when doing updates.
// This is not a problem on initial deployment since no resources have been created.
// DM is HIPAA compliant, so it's ok to leave its access.
// See https://cloud.google.com/iam/docs/granting-changing-revoking-access.
func grantDeploymentManagerAccess(project *config.Project) error {
	pnum := project.GeneratedFields.ProjectNumber
	if pnum == "" {
		return fmt.Errorf("project number not set in generated fields %+v", project.GeneratedFields)
	}
	serviceAcct := fmt.Sprintf("serviceAccount:%s@cloudservices.gserviceaccount.com", pnum)

	// TODO: account for this in the rule generator.
	for _, role := range deploymentManagerRoles {
		cmd := exec.Command(
			"gcloud", "projects", "add-iam-policy-binding", project.ID,
			"--role", role,
			"--member", serviceAcct,
			"--project", project.ID,
		)
		if err := cmdRun(cmd); err != nil {
			return fmt.Errorf("failed to grant role %q to DM service account %q: %v", role, serviceAcct, err)
		}
	}
	return nil
}

func getLogSinkServiceAccount(project *config.Project) (string, error) {
	cmd := exec.Command("gcloud", "logging", "sinks", "describe", project.BQLogSink.Name(), "--format", "json", "--project", project.ID)

	out, err := cmdOutput(cmd)
	if err != nil {
		return "", fmt.Errorf("failed to query log sink service account from gcloud: %v", err)
	}

	type sink struct {
		WriterIdentity string `json:"writerIdentity"`
	}

	s := new(sink)
	if err := json.Unmarshal(out, s); err != nil {
		return "", fmt.Errorf("failed to unmarshal sink output: %v", err)
	}
	return strings.TrimPrefix(s.WriterIdentity, "serviceAccount:"), nil
}

func deployAudit(project, auditProject *config.Project) error {
	rs := []config.Resource{&project.AuditLogs.LogsBQDataset}
	if project.AuditLogs.LogsGCSBucket != nil {
		rs = append(rs, project.AuditLogs.LogsGCSBucket)
	}
	deployment, err := getDeployment(project, rs)
	if err != nil {
		return err
	}

	// Append project ID to deployment name so each project has unique deployment if there is
	// a remote audit logs project.
	name := fmt.Sprintf("%s-%s", auditDeploymentName, project.ID)
	if err := upsertDeployment(name, deployment, auditProject.ID); err != nil {
		return fmt.Errorf("failed to deploy audit resources: %v", err)
	}
	return nil
}

func deployResources(project *config.Project) error {
	rs := project.DeploymentManagerResources()
	if len(rs) == 0 {
		log.Println("No resources to deploy.")
		return nil
	}
	deployment, err := getDeployment(project, rs)
	if err != nil {
		return err
	}
	if err := upsertDeployment(resourceDeploymentName, deployment, project.ID); err != nil {
		return fmt.Errorf("failed to deploy deployment manager resources: %v", err)
	}
	return nil
}

func getDeployment(project *config.Project, resources []config.Resource) (*deploymentmanager.Deployment, error) {
	deployment := &deploymentmanager.Deployment{}

	importSet := make(map[string]bool)

	for _, r := range resources {
		var typ string
		if typer, ok := r.(deploymentManagerTyper); ok {
			typ = typer.DeploymentManagerType()
		} else if pather, ok := r.(deploymentManagerPather); ok {
			var err error
			typ, err = filepath.Abs(pather.TemplatePath())
			if err != nil {
				return nil, fmt.Errorf("failed to get absolute path for %q: %v", pather.TemplatePath(), err)
			}
			if !importSet[typ] {
				deployment.Imports = append(deployment.Imports, &deploymentmanager.Import{Path: typ})
				importSet[typ] = true
			}
		} else {
			return nil, fmt.Errorf("failed to get type of %+v", r)
		}

		b, err := json.Marshal(r)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal resource: %v", err)
		}

		type resourceProperties struct {
			Properties map[string]interface{} `json:"properties"`
		}
		rp := new(resourceProperties)
		if err := json.Unmarshal(b, &rp); err != nil {
			return nil, fmt.Errorf("failed to unmarshal resource: %v", err)
		}

		res := &deploymentmanager.Resource{
			Name:       r.Name(),
			Type:       typ,
			Properties: rp.Properties,
		}

		if dr, ok := r.(depender); ok && len(dr.Dependencies()) > 0 {
			res.Metadata = &deploymentmanager.Metadata{DependsOn: dr.Dependencies()}
		}

		deployment.Resources = append(deployment.Resources, res)
	}

	return deployment, nil
}

func removeOwnerUser(project *config.Project) error {
	cmd := exec.Command("gcloud", "config", "get-value", "account", "--format", "json", "--project", project.ID)
	out, err := cmdOutput(cmd)
	if err != nil {
		return fmt.Errorf("failed to get currently authenticated user: %v", err)
	}
	var member string
	if err := json.Unmarshal(out, &member); err != nil {
		return fmt.Errorf("failed to unmarshal current user: %v", err)
	}
	role := "roles/owner"
	member = "user:" + member

	// TODO: check user specified bindings in case user wants the binding left
	has, err := hasBinding(project, role, member)
	if err != nil {
		return err
	}
	if !has {
		log.Printf("owner user %q already removed", member)
		return nil
	}

	cmd = exec.Command(
		"gcloud", "projects", "remove-iam-policy-binding", project.ID,
		"--member", member, "--role", role, "--project", project.ID)
	return cmdRun(cmd)
}

func hasBinding(project *config.Project, role string, member string) (has bool, err error) {
	cmd := exec.Command(
		"gcloud", "projects", "get-iam-policy", project.ID,
		"--project", project.ID,
		"--format", "json",
	)
	out, err := cmdOutput(cmd)
	if err != nil {
		return false, fmt.Errorf("failed to get iam policy bindings: %v", err)
	}
	log.Printf("Looking for role %q, member %q in:\n%v", role, member, string(out))

	type policy struct {
		Bindings []config.Binding `json:"bindings"`
	}
	p := new(policy)
	if err := json.Unmarshal(out, p); err != nil {
		return false, fmt.Errorf("failed to unmarshal get-iam-policy output: %v", err)
	}
	for _, b := range p.Bindings {
		if b.Role == role {
			for _, m := range b.Members {
				if m == member {
					return true, nil
				}
			}
		}
	}

	return false, nil
}

// deployPrerequisite deploys the CHC resources in the project.
func deployPrerequisite(project *config.Project) error {
	if err := upsertDeploymentFromFile(setupPrerequisiteDeploymentName, "deploy/config/templates/chc_resource/chc_res_type_provider.yaml", project.ID); err != nil {
		return fmt.Errorf("failed to deploy CHC type provider: %v", err)
	}
	return nil
}
<commit_msg>don't set stdout for output<commit_after>// Package apply provides utilities to apply a project config to GCP by deploying all defined resources.
package apply

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/healthcare/deploy/config"
	"github.com/GoogleCloudPlatform/healthcare/deploy/deploymentmanager"
)

const (
	deploymentNamePrefix            = "data-protect-toolkit"
	auditDeploymentName             = deploymentNamePrefix + "-audit"
	resourceDeploymentName          = deploymentNamePrefix + "-resources"
	setupPrerequisiteDeploymentName = deploymentNamePrefix + "-prerequisites"
)

// deploymentManagerRoles are the roles granted to the DM service account.
var deploymentManagerRoles = []string{"roles/owner", "roles/storage.admin"}

// deploymentRetryWaitTime is the time to wait between retrying a deployment to allow for concurrent operations to finish.
const deploymentRetryWaitTime = time.Minute

// The following vars are stubbed in tests.
var (
	cmdOutput = func(cmd *exec.Cmd) ([]byte, error) {
		log.Printf("Running: %v", cmd.Args)
		cmd.Stderr = os.Stderr
		return cmd.Output()
	}
	cmdRun = func(cmd *exec.Cmd) error {
		log.Printf("Running: %v", cmd.Args)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return cmd.Run()
	}
	upsertDeployment         = deploymentmanager.Upsert
	upsertDeploymentFromFile = deploymentmanager.UpsertFromFile
)

// parsedResource is an interface that must be implemented by all concrete resource implementations.
type parsedResource interface {
	Init(*config.Project) error
	Name() string
}

// deploymentManagerTyper should be implemented by resources that are natively supported by the deployment manager service.
// Use this if there is no suitable CFT template for a resource and a custom template is not needed.
// See https://cloud.google.com/deployment-manager/docs/configuration/supported-resource-types for valid types.
type deploymentManagerTyper interface {
	DeploymentManagerType() string
}

// deploymentManagerPather should be implemented by resources that use a DM template to deploy.
// Use this if the resource wraps a CFT or custom template.
type deploymentManagerPather interface {
	TemplatePath() string
}

// depender is the interface that defines a method to get dependent resources.
type depender interface {
	// Dependencies returns the name of the resource IDs to depend on.
	Dependencies() []string
}

// Apply deploys the CFT resources in the project.
func Apply(conf *config.Config, project *config.Project) error {
	if err := grantDeploymentManagerAccess(project); err != nil {
		return fmt.Errorf("failed to grant deployment manager access to the project: %v", err)
	}

	if err := deployPrerequisite(project); err != nil {
		return fmt.Errorf("failed to deploy CHC Type Provider: %v", err)
	}

	// TODO: stop retrying once
	// https://github.com/GoogleCloudPlatform/cloud-foundation-toolkit/issues/17
	// is fixed.
	for i := 0; i < 2; i++ {
		if err := deployResources(project); err == nil {
			break
		} else if i == 1 {
			return fmt.Errorf("failed to deploy resources: %v", err)
		}
		log.Printf("Sleeping for %v and retrying in case failure was due to concurrent IAM policy update", deploymentRetryWaitTime)
		time.Sleep(deploymentRetryWaitTime)
	}

	// Always get the latest log sink writer as when the sink is moved between deployments it may
	// create a new sink writer.
	sinkSA, err := getLogSinkServiceAccount(project)
	if err != nil {
		return fmt.Errorf("failed to get log sink service account: %v", err)
	}

	// Note: if the project was deployed the project Init will already have added the log sink service account permission on the dataset.
	if currSA := project.GeneratedFields.LogSinkServiceAccount; currSA == "" {
		project.AuditLogs.LogsBQDataset.Accesses = append(project.AuditLogs.LogsBQDataset.Accesses, &config.Access{
			Role: "WRITER", UserByEmail: sinkSA,
		})
	} else if currSA != sinkSA {
		project.GeneratedFields.LogSinkServiceAccount = sinkSA
		// Replace all instances of old writer SA with new.
		for _, a := range project.AuditLogs.LogsBQDataset.Accesses {
			if a.UserByEmail == currSA {
				a.UserByEmail = sinkSA
			}
		}
	}

	if err := deployAudit(project, conf.ProjectForAuditLogs(project)); err != nil {
		return fmt.Errorf("failed to deploy audit resources: %v", err)
	}

	if err := deployGKEWorkloads(project); err != nil {
		return fmt.Errorf("failed to deploy GKE workloads: %v", err)
	}

	// Only remove owner account if there is an organization to ensure the project has an administrator.
	if conf.Overall.OrganizationID != "" {
		if err := removeOwnerUser(project); err != nil {
			return fmt.Errorf("failed to remove owner user: %v", err)
		}
	}
	return nil
}

// grantDeploymentManagerAccess grants the necessary permissions to the DM service account to perform its actions.
// Note: we don't revoke deployment manager's access because permissions can take up to 7 minutes
// to propagate through the system, which can cause permission denied issues when doing updates.
// This is not a problem on initial deployment since no resources have been created.
// DM is HIPAA compliant, so it's ok to leave its access.
// See https://cloud.google.com/iam/docs/granting-changing-revoking-access.
func grantDeploymentManagerAccess(project *config.Project) error {
	pnum := project.GeneratedFields.ProjectNumber
	if pnum == "" {
		return fmt.Errorf("project number not set in generated fields %+v", project.GeneratedFields)
	}
	serviceAcct := fmt.Sprintf("serviceAccount:%s@cloudservices.gserviceaccount.com", pnum)

	// TODO: account for this in the rule generator.
	for _, role := range deploymentManagerRoles {
		cmd := exec.Command(
			"gcloud", "projects", "add-iam-policy-binding", project.ID,
			"--role", role,
			"--member", serviceAcct,
			"--project", project.ID,
		)
		if err := cmdRun(cmd); err != nil {
			return fmt.Errorf("failed to grant role %q to DM service account %q: %v", role, serviceAcct, err)
		}
	}
	return nil
}

func getLogSinkServiceAccount(project *config.Project) (string, error) {
	cmd := exec.Command("gcloud", "logging", "sinks", "describe", project.BQLogSink.Name(), "--format", "json", "--project", project.ID)

	out, err := cmdOutput(cmd)
	if err != nil {
		return "", fmt.Errorf("failed to query log sink service account from gcloud: %v", err)
	}

	type sink struct {
		WriterIdentity string `json:"writerIdentity"`
	}

	s := new(sink)
	if err := json.Unmarshal(out, s); err != nil {
		return "", fmt.Errorf("failed to unmarshal sink output: %v", err)
	}
	return strings.TrimPrefix(s.WriterIdentity, "serviceAccount:"), nil
}

func deployAudit(project, auditProject *config.Project) error {
	rs := []config.Resource{&project.AuditLogs.LogsBQDataset}
	if project.AuditLogs.LogsGCSBucket != nil {
		rs = append(rs, project.AuditLogs.LogsGCSBucket)
	}
	deployment, err := getDeployment(project, rs)
	if err != nil {
		return err
	}

	// Append project ID to deployment name so each project has unique deployment if there is
	// a remote audit logs project.
	name := fmt.Sprintf("%s-%s", auditDeploymentName, project.ID)
	if err := upsertDeployment(name, deployment, auditProject.ID); err != nil {
		return fmt.Errorf("failed to deploy audit resources: %v", err)
	}
	return nil
}

func deployResources(project *config.Project) error {
	rs := project.DeploymentManagerResources()
	if len(rs) == 0 {
		log.Println("No resources to deploy.")
		return nil
	}
	deployment, err := getDeployment(project, rs)
	if err != nil {
		return err
	}
	if err := upsertDeployment(resourceDeploymentName, deployment, project.ID); err != nil {
		return fmt.Errorf("failed to deploy deployment manager resources: %v", err)
	}
	return nil
}

func getDeployment(project *config.Project, resources []config.Resource) (*deploymentmanager.Deployment, error) {
	deployment := &deploymentmanager.Deployment{}

	importSet := make(map[string]bool)

	for _, r := range resources {
		var typ string
		if typer, ok := r.(deploymentManagerTyper); ok {
			typ = typer.DeploymentManagerType()
		} else if pather, ok := r.(deploymentManagerPather); ok {
			var err error
			typ, err = filepath.Abs(pather.TemplatePath())
			if err != nil {
				return nil, fmt.Errorf("failed to get absolute path for %q: %v", pather.TemplatePath(), err)
			}
			if !importSet[typ] {
				deployment.Imports = append(deployment.Imports, &deploymentmanager.Import{Path: typ})
				importSet[typ] = true
			}
		} else {
			return nil, fmt.Errorf("failed to get type of %+v", r)
		}

		b, err := json.Marshal(r)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal resource: %v", err)
		}

		type resourceProperties struct {
			Properties map[string]interface{} `json:"properties"`
		}
		rp := new(resourceProperties)
		if err := json.Unmarshal(b, &rp); err != nil {
			return nil, fmt.Errorf("failed to unmarshal resource: %v", err)
		}

		res := &deploymentmanager.Resource{
			Name:       r.Name(),
			Type:       typ,
			Properties: rp.Properties,
		}

		if dr, ok := r.(depender); ok && len(dr.Dependencies()) > 0 {
			res.Metadata = &deploymentmanager.Metadata{DependsOn: dr.Dependencies()}
		}

		deployment.Resources = append(deployment.Resources, res)
	}

	return deployment, nil
}

func removeOwnerUser(project *config.Project) error {
	cmd := exec.Command("gcloud", "config", "get-value", "account", "--format", "json", "--project", project.ID)
	out, err := cmdOutput(cmd)
	if err != nil {
		return fmt.Errorf("failed to get currently authenticated user: %v", err)
	}
	var member string
	if err := json.Unmarshal(out, &member); err != nil {
		return fmt.Errorf("failed to unmarshal current user: %v", err)
	}
	role := "roles/owner"
	member = "user:" + member

	// TODO: check user specified bindings in case user wants the binding left
	has, err := hasBinding(project, role, member)
	if err != nil {
		return err
	}
	if !has {
		log.Printf("owner user %q already removed", member)
		return nil
	}

	cmd = exec.Command(
		"gcloud", "projects", "remove-iam-policy-binding", project.ID,
		"--member", member, "--role", role, "--project", project.ID)
	return cmdRun(cmd)
}

func hasBinding(project *config.Project, role string, member string) (has bool, err error) {
	cmd := exec.Command(
		"gcloud", "projects", "get-iam-policy", project.ID,
		"--project", project.ID,
		"--format", "json",
	)
	out, err := cmdOutput(cmd)
	if err != nil {
		return false, fmt.Errorf("failed to get iam policy bindings: %v", err)
	}
	log.Printf("Looking for role %q, member %q in:\n%v", role, member, string(out))

	type policy struct {
		Bindings []config.Binding `json:"bindings"`
	}
	p := new(policy)
	if err := json.Unmarshal(out, p); err != nil {
		return false, fmt.Errorf("failed to unmarshal get-iam-policy output: %v", err)
	}
	for _, b := range p.Bindings {
		if b.Role == role {
			for _, m := range b.Members {
				if m == member {
					return true, nil
				}
			}
		}
	}

	return false, nil
}

// deployPrerequisite deploys the CHC resources in the project.
func deployPrerequisite(project *config.Project) error {
	if err := upsertDeploymentFromFile(setupPrerequisiteDeploymentName, "deploy/config/templates/chc_resource/chc_res_type_provider.yaml", project.ID); err != nil {
		return fmt.Errorf("failed to deploy CHC type provider: %v", err)
	}
	return nil
}
<|endoftext|>


<commit_before>package attr

// Getter is an interface which is implemented by types which want to export typed
// values.
type Getter interface {
	GetInt(string) int
	GetString(string) string
	GetStrings(string) []string
}

// Interface is a type which defines behaviour necessary to implement a typed attribute.
type Interface interface {
	Field() string
	IsEmpty(x interface{}) bool
	Value(Getter) interface{}
	Intersect(x, y interface{}) interface{}
}

type valueType struct {
	field string
	empty interface{}
	get   func(Getter) interface{}
}

func (v *valueType) Field() string {
	return v.field
}

func (v *valueType) IsEmpty(x interface{}) bool {
	return v.empty == x
}

func (v *valueType) Value(g Getter) interface{} {
	return v.get(g)
}

func (v *valueType) Intersect(x, y interface{}) interface{} {
	if x == y {
		return x
	}
	return v.empty
}

func String(f string) Interface {
	return &valueType{
		field: f,
		empty: "",
		get: func(g Getter) interface{} {
			return g.GetString(f)
		},
	}
}

func Int(f string) Interface {
	return &valueType{
		field: f,
		empty: 0,
		get: func(g Getter) interface{} {
			return g.GetInt(f)
		},
	}
}

type stringsType struct {
	valueType
}

func (p *stringsType) IsEmpty(x interface{}) bool {
	if x == nil {
		return true
	}
	xs := x.([]string)
	return len(xs) == 0
}

func (p *stringsType) Intersect(x, y interface{}) interface{} {
	if x == nil || y == nil {
		return nil
	}
	xs := x.([]string)
	ys := y.([]string)
	return stringSliceIntersect(xs, ys)
}

func Strings(f string) Interface {
	return &stringsType{
		valueType{
			field: f,
			empty: nil,
			get: func(g Getter) interface{} {
				return g.GetStrings(f)
			},
		},
	}
}

// stringSliceIntersect computes the intersection of two string slices (ignoring ordering).
func stringSliceIntersect(s, t []string) []string {
	var res []string
	m := make(map[string]bool)
	for _, x := range s {
		m[x] = true
	}
	for _, y := range t {
		if m[y] {
			res = append(res, y)
		}
	}
	return res
}
<commit_msg>Add license comment to attr.go.<commit_after>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package attr

// Getter is an interface which is implemented by types which want to export typed
// values.
type Getter interface {
	GetInt(string) int
	GetString(string) string
	GetStrings(string) []string
}

// Interface is a type which defines behaviour necessary to implement a typed attribute.
type Interface interface {
	Field() string
	IsEmpty(x interface{}) bool
	Value(Getter) interface{}
	Intersect(x, y interface{}) interface{}
}

type valueType struct {
	field string
	empty interface{}
	get   func(Getter) interface{}
}

func (v *valueType) Field() string {
	return v.field
}

func (v *valueType) IsEmpty(x interface{}) bool {
	return v.empty == x
}

func (v *valueType) Value(g Getter) interface{} {
	return v.get(g)
}

func (v *valueType) Intersect(x, y interface{}) interface{} {
	if x == y {
		return x
	}
	return v.empty
}

func String(f string) Interface {
	return &valueType{
		field: f,
		empty: "",
		get: func(g Getter) interface{} {
			return g.GetString(f)
		},
	}
}

func Int(f string) Interface {
	return &valueType{
		field: f,
		empty: 0,
		get: func(g Getter) interface{} {
			return g.GetInt(f)
		},
	}
}

type stringsType struct {
	valueType
}

func (p *stringsType) IsEmpty(x interface{}) bool {
	if x == nil {
		return true
	}
	xs := x.([]string)
	return len(xs) == 0
}

func (p *stringsType) Intersect(x, y interface{}) interface{} {
	if x == nil || y == nil {
		return nil
	}
	xs := x.([]string)
	ys := y.([]string)
	return stringSliceIntersect(xs, ys)
}

func Strings(f string) Interface {
	return &stringsType{
		valueType{
			field: f,
			empty: nil,
			get: func(g Getter) interface{} {
				return g.GetStrings(f)
			},
		},
	}
}

// stringSliceIntersect computes the intersection of two string slices (ignoring ordering).
func stringSliceIntersect(s, t []string) []string {
	var res []string
	m := make(map[string]bool)
	for _, x := range s {
		m[x] = true
	}
	for _, y := range t {
		if m[y] {
			res = append(res, y)
		}
	}
	return res
}
<|endoftext|>


<commit_before>package containerd

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/containerd/containerd"
	"github.com/containerd/containerd/api/services/execution"
	"github.com/containerd/containerd/api/types/task"
	dockermount "github.com/docker/docker/pkg/mount"
	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/api/naming"
	"github.com/docker/swarmkit/log"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

var (
	devNull                    *os.File
	errAdapterNotPrepared      = errors.New("container adapter not prepared")
	mountPropagationReverseMap = map[api.Mount_BindOptions_MountPropagation]string{
		api.MountPropagationPrivate:  "private",
		api.MountPropagationRPrivate: "rprivate",
		api.MountPropagationShared:   "shared",
		api.MountPropagationRShared:  "rshared",
		api.MountPropagationRSlave:   "slave",
		api.MountPropagationSlave:    "rslave",
	}
)

// containerAdapter conducts remote operations for a container. All calls
// are mostly naked calls to the client API, seeded with information from
// containerConfig.
type containerAdapter struct {
	client         *containerd.Client
	spec           *api.ContainerSpec
	secrets        exec.SecretGetter
	name           string
	image          containerd.Image // Pulled image
	container      containerd.Container
	task           containerd.Task
	deleteResponse *execution.DeleteResponse
}

func newContainerAdapter(client *containerd.Client, task *api.Task, secrets exec.SecretGetter) (*containerAdapter, error) {
	spec := task.Spec.GetContainer()
	if spec == nil {
		return nil, exec.ErrRuntimeUnsupported
	}

	c := &containerAdapter{
		client:  client,
		spec:    spec,
		secrets: secrets,
		name:    naming.Task(task),
	}

	if err := c.reattach(context.Background()); err != nil {
		return nil, err
	}

	return c, nil
}

// reattaches to an existing container. If the container is found but
// the task is missing then still succeeds, allowing subsequent use of
// c.delete()
func (c *containerAdapter) reattach(ctx context.Context) error {
	container, err := c.client.LoadContainer(ctx, c.name)
	if err != nil {
		if grpc.Code(err) == codes.NotFound {
			c.log(ctx).Debug("reattach: container not found")
			return nil
		}

		return errors.Wrap(err, "reattach: loading container")
	}
	c.log(ctx).Debug("reattach: loaded container")
	c.container = container

	// TODO(ijc) Consider an addition to container library which
	// directly attaches stdin to /dev/null.
	if devNull == nil {
		if devNull, err = os.Open(os.DevNull); err != nil {
			return errors.Wrap(err, "reattach: opening null device")
		}
	}

	task, err := container.Task(ctx, containerd.WithAttach(devNull, os.Stdout, os.Stderr))
	if err != nil {
		if err == containerd.ErrNoRunningTask {
			c.log(ctx).WithError(err).Info("reattach: no running task")
			return nil
		}
		return errors.Wrap(err, "reattach: reattaching task")
	}
	c.task = task
	c.log(ctx).Debug("reattach: successful")
	return nil
}

func (c *containerAdapter) log(ctx context.Context) *logrus.Entry {
	return log.G(ctx).WithFields(logrus.Fields{
		"ID": c.name,
	})
}

func (c *containerAdapter) pullImage(ctx context.Context) error {
	image, err := c.client.Pull(ctx, c.spec.Image, containerd.WithPullUnpack)
	if err != nil {
		return errors.Wrap(err, "pulling container image")
	}
	c.image = image

	return nil
}

func withMounts(ctx context.Context, ms []api.Mount) containerd.SpecOpts {
	sort.Sort(mounts(ms))

	return func(s *specs.Spec) error {
		for _, m := range ms {
			if !filepath.IsAbs(m.Target) {
				return errors.Errorf("mount %s is not absolute", m.Target)
			}

			switch m.Type {
			case api.MountTypeTmpfs:
				opts := []string{"noexec", "nosuid", "nodev", "rprivate"}
				if m.TmpfsOptions != nil {
					if m.TmpfsOptions.SizeBytes <= 0 {
						return errors.New("invalid tmpfs size give")
					}
					opts = append(opts, fmt.Sprintf("size=%d", m.TmpfsOptions.SizeBytes))
					opts = append(opts, fmt.Sprintf("mode=%o", m.TmpfsOptions.Mode))
				}
				if m.ReadOnly {
					opts = append(opts, "ro")
				} else {
					opts = append(opts, "rw")
				}

				opts, err := dockermount.MergeTmpfsOptions(opts)
				if err != nil {
					return err
				}

				s.Mounts = append(s.Mounts, specs.Mount{
					Destination: m.Target,
					Type:        "tmpfs",
					Source:      "tmpfs",
					Options:     opts,
				})

			case api.MountTypeVolume:
				return errors.Errorf("volume mounts not implemented, ignoring %v", m)

			case api.MountTypeBind:
				opts := []string{"rbind"}
				if m.ReadOnly {
					opts = append(opts, "ro")
				} else {
					opts = append(opts, "rw")
				}

				propagation := "rprivate"
				if m.BindOptions != nil {
					if p, ok := mountPropagationReverseMap[m.BindOptions.Propagation]; ok {
						propagation = p
					} else {
						log.G(ctx).Warningf("unknown bind mount propagation, using %q", propagation)
					}
				}
				opts = append(opts, propagation)

				s.Mounts = append(s.Mounts, specs.Mount{
					Destination: m.Target,
					Type:        "bind",
					Source:      m.Source,
					Options:     opts,
				})
			}
		}
		return nil
	}
}

func (c *containerAdapter) isPrepared() bool {
	return c.container != nil && c.task != nil
}

func (c *containerAdapter) prepare(ctx context.Context) error {
	if c.isPrepared() {
		return errors.New("adapter already prepared")
	}
	if c.image == nil {
		return errors.New("image has not been pulled")
	}

	specOpts := []containerd.SpecOpts{
		containerd.WithImageConfig(ctx, c.image),
		withMounts(ctx, c.spec.Mounts),
	}

	// spec.Process.Args is config.Entrypoint + config.Cmd at this
	// point from WithImageConfig above. If the ContainerSpec
	// specifies a Command then we can completely override. If it
	// does not then all we can do is append our Args and hope
	// they do not conflict.
	// TODO(ijc) Improve this
	if len(c.spec.Command) > 0 {
		args := append(c.spec.Command, c.spec.Args...)
		specOpts = append(specOpts, containerd.WithProcessArgs(args...))
	} else {
		specOpts = append(specOpts, func(s *specs.Spec) error {
			s.Process.Args = append(s.Process.Args, c.spec.Args...)
			return nil
		})
	}

	spec, err := containerd.GenerateSpec(specOpts...)
	if err != nil {
		return err
	}

	// TODO(ijc) Consider an addition to container library which
	// directly attaches stdin to /dev/null.
	if devNull == nil {
		if devNull, err = os.Open(os.DevNull); err != nil {
			return errors.Wrap(err, "opening null device")
		}
	}

	c.container, err = c.client.NewContainer(ctx, c.name,
		containerd.WithSpec(spec),
		containerd.WithNewRootFS(c.name, c.image))
	if err != nil {
		return errors.Wrap(err, "creating container")
	}

	// TODO(ijc) support ControllerLogs interface.
	io := containerd.NewIOWithTerminal(devNull, os.Stdout, os.Stderr, spec.Process.Terminal)

	c.task, err = c.container.NewTask(ctx, io)
	if err != nil {
		// Destroy the container we created above, but
		// propagate the original error.
		if err2 := c.container.Delete(ctx); err2 != nil {
			c.log(ctx).WithError(err2).Error("failed to delete container on prepare failure")
		}
		c.container = nil
		return errors.Wrap(err, "creating task")
	}

	return nil
}

func (c *containerAdapter) start(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	tasks := c.client.TaskService()

	_, err := tasks.Start(ctx, &execution.StartRequest{
		ContainerID: c.name,
	})
	return err
}

func (c *containerAdapter) eventStream(ctx context.Context, id string) (<-chan task.Event, <-chan error, error) {

	var (
		evtch = make(chan task.Event)
		errch = make(chan error)
	)

	return evtch, errch, nil
}

// events issues a call to the events API and returns a channel with all
// events. The stream of events can be shutdown by cancelling the context.
//
// A chan struct{} is returned that will be closed if the event processing
// fails and needs to be restarted.
func (c *containerAdapter) events(ctx context.Context, opts ...grpc.CallOption) (<-chan task.Event, <-chan struct{}, error) {
	if !c.isPrepared() {
		return nil, nil, errAdapterNotPrepared
	}

	// TODO(stevvooe): Move this to a single, global event dispatch. For
	// now, we create a connection per container.
	var (
		eventsq = make(chan task.Event)
		closed  = make(chan struct{})
	)

	c.log(ctx).Debugf("waiting on events")

	tasks := c.client.TaskService()
	cl, err := tasks.Events(ctx, &execution.EventsRequest{}, opts...)
	if err != nil {
		c.log(ctx).WithError(err).Errorf("failed to start event stream")
		return nil, nil, err
	}

	go func() {
		defer close(closed)

		for {
			evt, err := cl.Recv()
			if err != nil {
				c.log(ctx).WithError(err).Error("fatal error from events stream")
				return
			}
			if evt.ID != c.name {
				c.log(ctx).Debugf("Event for a different container %s", evt.ID)
				continue
			}

			select {
			case eventsq <- *evt:
			case <-ctx.Done():
				return
			}
		}
	}()

	return eventsq, closed, nil
}

func (c *containerAdapter) inspect(ctx context.Context) (task.Task, error) {
	if !c.isPrepared() {
		return task.Task{}, errAdapterNotPrepared
	}

	tasks := c.client.TaskService()
	rsp, err := tasks.Info(ctx, &execution.InfoRequest{ContainerID: c.name})
	if err != nil {
		return task.Task{}, err
	}
	return *rsp.Task, nil
}

func (c *containerAdapter) shutdown(ctx context.Context) (uint32, error) {
	if !c.isPrepared() {
		return 0, errAdapterNotPrepared
	}

	if c.deleteResponse == nil {
		var err error
		c.log(ctx).Debug("Deleting")

		tasks := c.client.TaskService()
		rsp, err := tasks.Delete(ctx, &execution.DeleteRequest{ContainerID: c.name})
		if err != nil {
			return 0, err
		}
		c.log(ctx).Debugf("Status=%d", rsp.ExitStatus)
		c.deleteResponse = rsp

		// TODO(ijc) this should be moved to the remove method.
		if err := c.container.Delete(ctx); err != nil {
			c.log(ctx).WithError(err).Warnf("failed to delete container")
		}
	}

	return c.deleteResponse.ExitStatus, nil
}

func (c *containerAdapter) terminate(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	c.log(ctx).Debug("Terminate")
	return errors.New("terminate not implemented")
}

func (c *containerAdapter) remove(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	c.log(ctx).Debug("Remove")
	return nil
}

func isContainerCreateNameConflict(err error) bool {
	// container ".*" already exists
	splits := strings.SplitN(err.Error(), "\"", 3)
	return splits[0] == "container " && splits[2] == " already exists"
}

func isUnknownContainer(err error) bool {
	return strings.Contains(err.Error(), "container does not exist")
}

// For sort.Sort
type mounts []api.Mount

// Len returns the number of mounts. Used in sorting.
func (m mounts) Len() int {
	return len(m)
}

// Less returns true if the number of parts (a/b/c would be 3 parts) in the
// mount indexed by parameter 1 is less than that of the mount indexed by
// parameter 2. Used in sorting.
func (m mounts) Less(i, j int) bool {
	return m.parts(i) < m.parts(j)
}

// Swap swaps two items in an array of mounts. Used in sorting
func (m mounts) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

// parts returns the number of parts in the destination of a mount. Used in sorting.
func (m mounts) parts(i int) int {
	return strings.Count(filepath.Clean(m[i].Target), string(os.PathSeparator))
}
<commit_msg>containerd: Use task.Start instead of open coding<commit_after>package containerd

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/containerd/containerd"
	"github.com/containerd/containerd/api/services/execution"
	"github.com/containerd/containerd/api/types/task"
	dockermount "github.com/docker/docker/pkg/mount"
	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/api/naming"
	"github.com/docker/swarmkit/log"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

var (
	devNull                    *os.File
	errAdapterNotPrepared      = errors.New("container adapter not prepared")
	mountPropagationReverseMap = map[api.Mount_BindOptions_MountPropagation]string{
		api.MountPropagationPrivate:  "private",
		api.MountPropagationRPrivate: "rprivate",
		api.MountPropagationShared:   "shared",
		api.MountPropagationRShared:  "rshared",
		api.MountPropagationRSlave:   "slave",
		api.MountPropagationSlave:    "rslave",
	}
)

// containerAdapter conducts remote operations for a container. All calls
// are mostly naked calls to the client API, seeded with information from
// containerConfig.
type containerAdapter struct {
	client         *containerd.Client
	spec           *api.ContainerSpec
	secrets        exec.SecretGetter
	name           string
	image          containerd.Image // Pulled image
	container      containerd.Container
	task           containerd.Task
	deleteResponse *execution.DeleteResponse
}

func newContainerAdapter(client *containerd.Client, task *api.Task, secrets exec.SecretGetter) (*containerAdapter, error) {
	spec := task.Spec.GetContainer()
	if spec == nil {
		return nil, exec.ErrRuntimeUnsupported
	}

	c := &containerAdapter{
		client:  client,
		spec:    spec,
		secrets: secrets,
		name:    naming.Task(task),
	}

	if err := c.reattach(context.Background()); err != nil {
		return nil, err
	}

	return c, nil
}

// reattaches to an existing container. If the container is found but
// the task is missing then still succeeds, allowing subsequent use of
// c.delete()
func (c *containerAdapter) reattach(ctx context.Context) error {
	container, err := c.client.LoadContainer(ctx, c.name)
	if err != nil {
		if grpc.Code(err) == codes.NotFound {
			c.log(ctx).Debug("reattach: container not found")
			return nil
		}

		return errors.Wrap(err, "reattach: loading container")
	}
	c.log(ctx).Debug("reattach: loaded container")
	c.container = container

	// TODO(ijc) Consider an addition to container library which
	// directly attaches stdin to /dev/null.
	if devNull == nil {
		if devNull, err = os.Open(os.DevNull); err != nil {
			return errors.Wrap(err, "reattach: opening null device")
		}
	}

	task, err := container.Task(ctx, containerd.WithAttach(devNull, os.Stdout, os.Stderr))
	if err != nil {
		if err == containerd.ErrNoRunningTask {
			c.log(ctx).WithError(err).Info("reattach: no running task")
			return nil
		}
		return errors.Wrap(err, "reattach: reattaching task")
	}
	c.task = task
	c.log(ctx).Debug("reattach: successful")
	return nil
}

func (c *containerAdapter) log(ctx context.Context) *logrus.Entry {
	return log.G(ctx).WithFields(logrus.Fields{
		"ID": c.name,
	})
}

func (c *containerAdapter) pullImage(ctx context.Context) error {
	image, err := c.client.Pull(ctx, c.spec.Image, containerd.WithPullUnpack)
	if err != nil {
		return errors.Wrap(err, "pulling container image")
	}
	c.image = image

	return nil
}

func withMounts(ctx context.Context, ms []api.Mount) containerd.SpecOpts {
	sort.Sort(mounts(ms))

	return func(s *specs.Spec) error {
		for _, m := range ms {
			if !filepath.IsAbs(m.Target) {
				return errors.Errorf("mount %s is not absolute", m.Target)
			}

			switch m.Type {
			case api.MountTypeTmpfs:
				opts := []string{"noexec", "nosuid", "nodev", "rprivate"}
				if m.TmpfsOptions != nil {
					if m.TmpfsOptions.SizeBytes <= 0 {
						return errors.New("invalid tmpfs size give")
					}
					opts = append(opts, fmt.Sprintf("size=%d", m.TmpfsOptions.SizeBytes))
					opts = append(opts, fmt.Sprintf("mode=%o", m.TmpfsOptions.Mode))
				}
				if m.ReadOnly {
					opts = append(opts, "ro")
				} else {
					opts = append(opts, "rw")
				}

				opts, err := dockermount.MergeTmpfsOptions(opts)
				if err != nil {
					return err
				}

				s.Mounts = append(s.Mounts, specs.Mount{
					Destination: m.Target,
					Type:        "tmpfs",
					Source:      "tmpfs",
					Options:     opts,
				})

			case api.MountTypeVolume:
				return errors.Errorf("volume mounts not implemented, ignoring %v", m)

			case api.MountTypeBind:
				opts := []string{"rbind"}
				if m.ReadOnly {
					opts = append(opts, "ro")
				} else {
					opts = append(opts, "rw")
				}

				propagation := "rprivate"
				if m.BindOptions != nil {
					if p, ok := mountPropagationReverseMap[m.BindOptions.Propagation]; ok {
						propagation = p
					} else {
						log.G(ctx).Warningf("unknown bind mount propagation, using %q", propagation)
					}
				}
				opts = append(opts, propagation)

				s.Mounts = append(s.Mounts, specs.Mount{
					Destination: m.Target,
					Type:        "bind",
					Source:      m.Source,
					Options:     opts,
				})
			}
		}
		return nil
	}
}

func (c *containerAdapter) isPrepared() bool {
	return c.container != nil && c.task != nil
}

func (c *containerAdapter) prepare(ctx context.Context) error {
	if c.isPrepared() {
		return errors.New("adapter already prepared")
	}
	if c.image == nil {
		return errors.New("image has not been pulled")
	}

	specOpts := []containerd.SpecOpts{
		containerd.WithImageConfig(ctx, c.image),
		withMounts(ctx, c.spec.Mounts),
	}

	// spec.Process.Args is config.Entrypoint + config.Cmd at this
	// point from WithImageConfig above. If the ContainerSpec
	// specifies a Command then we can completely override. If it
	// does not then all we can do is append our Args and hope
	// they do not conflict.
	// TODO(ijc) Improve this
	if len(c.spec.Command) > 0 {
		args := append(c.spec.Command, c.spec.Args...)
		specOpts = append(specOpts, containerd.WithProcessArgs(args...))
	} else {
		specOpts = append(specOpts, func(s *specs.Spec) error {
			s.Process.Args = append(s.Process.Args, c.spec.Args...)
			return nil
		})
	}

	spec, err := containerd.GenerateSpec(specOpts...)
	if err != nil {
		return err
	}

	// TODO(ijc) Consider an addition to container library which
	// directly attaches stdin to /dev/null.
	if devNull == nil {
		if devNull, err = os.Open(os.DevNull); err != nil {
			return errors.Wrap(err, "opening null device")
		}
	}

	c.container, err = c.client.NewContainer(ctx, c.name,
		containerd.WithSpec(spec),
		containerd.WithNewRootFS(c.name, c.image))
	if err != nil {
		return errors.Wrap(err, "creating container")
	}

	// TODO(ijc) support ControllerLogs interface.
	io := containerd.NewIOWithTerminal(devNull, os.Stdout, os.Stderr, spec.Process.Terminal)

	c.task, err = c.container.NewTask(ctx, io)
	if err != nil {
		// Destroy the container we created above, but
		// propagate the original error.
		if err2 := c.container.Delete(ctx); err2 != nil {
			c.log(ctx).WithError(err2).Error("failed to delete container on prepare failure")
		}
		c.container = nil
		return errors.Wrap(err, "creating task")
	}

	return nil
}

func (c *containerAdapter) start(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}
	err := c.task.Start(ctx)
	return errors.Wrap(err, "starting")
}

func (c *containerAdapter) eventStream(ctx context.Context, id string) (<-chan task.Event, <-chan error, error) {

	var (
		evtch = make(chan task.Event)
		errch = make(chan error)
	)

	return evtch, errch, nil
}

// events issues a call to the events API and returns a channel with all
// events. The stream of events can be shutdown by cancelling the context.
//
// A chan struct{} is returned that will be closed if the event processing
// fails and needs to be restarted.
func (c *containerAdapter) events(ctx context.Context, opts ...grpc.CallOption) (<-chan task.Event, <-chan struct{}, error) {
	if !c.isPrepared() {
		return nil, nil, errAdapterNotPrepared
	}

	// TODO(stevvooe): Move this to a single, global event dispatch. For
	// now, we create a connection per container.
	var (
		eventsq = make(chan task.Event)
		closed  = make(chan struct{})
	)

	c.log(ctx).Debugf("waiting on events")

	tasks := c.client.TaskService()
	cl, err := tasks.Events(ctx, &execution.EventsRequest{}, opts...)
	if err != nil {
		c.log(ctx).WithError(err).Errorf("failed to start event stream")
		return nil, nil, err
	}

	go func() {
		defer close(closed)

		for {
			evt, err := cl.Recv()
			if err != nil {
				c.log(ctx).WithError(err).Error("fatal error from events stream")
				return
			}
			if evt.ID != c.name {
				c.log(ctx).Debugf("Event for a different container %s", evt.ID)
				continue
			}

			select {
			case eventsq <- *evt:
			case <-ctx.Done():
				return
			}
		}
	}()

	return eventsq, closed, nil
}

func (c *containerAdapter) inspect(ctx context.Context) (task.Task, error) {
	if !c.isPrepared() {
		return task.Task{}, errAdapterNotPrepared
	}

	tasks := c.client.TaskService()
	rsp, err := tasks.Info(ctx, &execution.InfoRequest{ContainerID: c.name})
	if err != nil {
		return task.Task{}, err
	}
	return *rsp.Task, nil
}

func (c *containerAdapter) shutdown(ctx context.Context) (uint32, error) {
	if !c.isPrepared() {
		return 0, errAdapterNotPrepared
	}

	if c.deleteResponse == nil {
		var err error
		c.log(ctx).Debug("Deleting")

		tasks := c.client.TaskService()
		rsp, err := tasks.Delete(ctx, &execution.DeleteRequest{ContainerID: c.name})
		if err != nil {
			return 0, err
		}
		c.log(ctx).Debugf("Status=%d", rsp.ExitStatus)
		c.deleteResponse = rsp

		// TODO(ijc) this should be moved to the remove method.
		if err := c.container.Delete(ctx); err != nil {
			c.log(ctx).WithError(err).Warnf("failed to delete container")
		}
	}

	return c.deleteResponse.ExitStatus, nil
}

func (c *containerAdapter) terminate(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	c.log(ctx).Debug("Terminate")
	return errors.New("terminate not implemented")
}

func (c *containerAdapter) remove(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	c.log(ctx).Debug("Remove")
	return nil
}

func isContainerCreateNameConflict(err error) bool {
	// container ".*" already exists
	splits := strings.SplitN(err.Error(), "\"", 3)
	return splits[0] == "container " && splits[2] == " already exists"
}

func isUnknownContainer(err error) bool {
	return strings.Contains(err.Error(), "container does not exist")
}

// For sort.Sort
type mounts []api.Mount

// Len returns the number of mounts. Used in sorting.
func (m mounts) Len() int {
	return len(m)
}

// Less returns true if the number of parts (a/b/c would be 3 parts) in the
// mount indexed by parameter 1 is less than that of the mount indexed by
// parameter 2. Used in sorting.
func (m mounts) Less(i, j int) bool {
	return m.parts(i) < m.parts(j)
}

// Swap swaps two items in an array of mounts. Used in sorting
func (m mounts) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

// parts returns the number of parts in the destination of a mount. Used in sorting.
func (m mounts) parts(i int) int {
	return strings.Count(filepath.Clean(m[i].Target), string(os.PathSeparator))
}
<|endoftext|>


<commit_before>package elb

import (
	"fmt"

	"errors"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	aws_elb "github.com/aws/aws-sdk-go/service/elb"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sky-uk/feed/ingress"
)

const (
	elbTag = "sky.uk/KubernetesClusterFrontend"
)

var attachedFrontendGauge = prometheus.NewGauge(prometheus.GaugeOpts{
	Namespace: "feed",
	Subsystem: "ingress",
	Name:      "frontends_attached",
	Help:      "The total number of frontends attached",
})

// New  creates a new ELB frontend
func New(region string, clusterName string, expectedFrontends int) ingress.Frontend {
	log.Infof("ELB Front end region: %s cluster: %s expected frontends: %d", region, clusterName, expectedFrontends)
	metadata := ec2metadata.New(session.New())
	return &elb{
		metadata:          metadata,
		awsElb:            aws_elb.New(session.New(&aws.Config{Region: &region})),
		clusterName:       clusterName,
		region:            region,
		expectedFrontends: expectedFrontends,
		maxTagQuery:       20,
	}
}

type elb struct {
	awsElb              ELB
	metadata            EC2Metadata
	clusterName         string
	region              string
	expectedFrontends   int
	maxTagQuery         int
	instanceID          string
	elbs                []string
	registeredFrontends int
}

// ELB interface to allow mocking of real calls to AWS as well as cutting down the methods from the real
// interface to only the ones we use
type ELB interface {
	DescribeLoadBalancers(input *aws_elb.DescribeLoadBalancersInput) (*aws_elb.DescribeLoadBalancersOutput, error)
	DescribeTags(input *aws_elb.DescribeTagsInput) (*aws_elb.DescribeTagsOutput, error)
	RegisterInstancesWithLoadBalancer(input *aws_elb.RegisterInstancesWithLoadBalancerInput) (*aws_elb.RegisterInstancesWithLoadBalancerOutput, error)
	DeregisterInstancesFromLoadBalancer(input *aws_elb.DeregisterInstancesFromLoadBalancerInput) (*aws_elb.DeregisterInstancesFromLoadBalancerOutput, error)
}

// EC2Metadata interface to allow mocking of the real calls to AWS
type EC2Metadata interface {
	Available() bool
	Region() (string, error)
	GetInstanceIdentityDocument() (ec2metadata.EC2InstanceIdentityDocument, error)
}

func (e *elb) Attach() error {

	if e.expectedFrontends == 0 {
		return nil
	}

	id, err := e.metadata.GetInstanceIdentityDocument()
	if err != nil {
		return fmt.Errorf("unable to query ec2 metadata service for InstanceId: %v", err)
	}

	instance := id.InstanceID
	log.Infof("Attaching to ELBs from instance %s", instance)
	clusterFrontEnds, err := e.findFrontEndElbs()

	if err != nil {
		return err
	}

	log.Infof("Found %d front ends", len(clusterFrontEnds))

	// Save these now so we can always know what we might have done
	// up until this point we have only read data
	e.elbs = clusterFrontEnds
	e.instanceID = instance
	registered := 0

	for _, frontend := range clusterFrontEnds {
		log.Infof("Registering instance %s with elb %", instance, frontend)
		_, err = e.awsElb.RegisterInstancesWithLoadBalancer(&aws_elb.RegisterInstancesWithLoadBalancerInput{
			Instances: []*aws_elb.Instance{
				&aws_elb.Instance{
					InstanceId: aws.String(instance),
				}},
			LoadBalancerName: aws.String(frontend),
		})

		if err != nil {
			return fmt.Errorf("unable to register instance %s with elb %s: %v", instance, frontend, err)
		}
		registered++

	}

	prometheus.Register(attachedFrontendGauge)
	attachedFrontendGauge.Set(float64(registered))
	e.registeredFrontends = registered
	return nil
}

func (e *elb) findFrontEndElbs() ([]string, error) {
	// Find the load balancers that are tagged with this cluster name
	request := &aws_elb.DescribeLoadBalancersInput{}
	var lbNames []*string
	for {
		resp, err := e.awsElb.DescribeLoadBalancers(request)

		if err != nil {
			return nil, fmt.Errorf("unable to describe load balancers: %v", err)
		}

		for _, entry := range resp.LoadBalancerDescriptions {
			lbNames = append(lbNames, entry.LoadBalancerName)
		}

		if resp.NextMarker == nil {
			break
		}

		// Set the next marker
		request = &aws_elb.DescribeLoadBalancersInput{
			Marker: resp.NextMarker,
		}
	}

	log.Infof("Found %d loadbalancers. Checking for %s tag set to %s", len(lbNames), elbTag, e.clusterName)
	var clusterFrontEnds []string
	totalLbs := len(lbNames)
	for i := 0; i < len(lbNames); i += e.maxTagQuery {
		to := min(i+e.maxTagQuery, totalLbs)
		log.Debugf("Querying tags from %d to %d", i, to)
		names := lbNames[i:to]
		output, err := e.awsElb.DescribeTags(&aws_elb.DescribeTagsInput{
			LoadBalancerNames: names,
		})

		if err != nil {
			return nil, fmt.Errorf("unable to describe tags: %v", err)
		}

		for _, description := range output.TagDescriptions {
			for _, tag := range description.Tags {
				if *tag.Key == elbTag && *tag.Value == e.clusterName {
					log.Infof("Found frontend elb %s", *description.LoadBalancerName)
					clusterFrontEnds = append(clusterFrontEnds, *description.LoadBalancerName)
				}
			}
		}
	}
	return clusterFrontEnds, nil
}

// Detach removes this instance from all the front end ELBs
func (e *elb) Detach() error {
	var failed = false
	for _, elb := range e.elbs {
		log.Infof("Deregistering instance %s with elb %s", e.instanceID, elb)
		_, err := e.awsElb.DeregisterInstancesFromLoadBalancer(&aws_elb.DeregisterInstancesFromLoadBalancerInput{
			Instances:        []*aws_elb.Instance{&aws_elb.Instance{InstanceId: aws.String(e.instanceID)}},
			LoadBalancerName: aws.String(elb),
		})

		if err != nil {
			log.Warnf("unable to deregister instance %s with elb %s: %v", e.instanceID, elb, err)
			failed = true
		}
	}
	if failed {
		return errors.New("at least one ELB failed to detach")
	}
	return nil
}

func (e *elb) Health() error {
	if e.registeredFrontends != e.expectedFrontends {
		return fmt.Errorf("expected frontends %d registered frontends %d", e.expectedFrontends, e.registeredFrontends)
	}
	return nil
}

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}
<commit_msg>Add synchronization around elb front end attach and health methods<commit_after>package elb

import (
	"fmt"

	"errors"

	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	aws_elb "github.com/aws/aws-sdk-go/service/elb"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sky-uk/feed/ingress"
)

const (
	elbTag = "sky.uk/KubernetesClusterFrontend"
)

var attachedFrontendGauge = prometheus.NewGauge(prometheus.GaugeOpts{
	Namespace: "feed",
	Subsystem: "ingress",
	Name:      "frontends_attached",
	Help:      "The total number of frontends attached",
})

// New  creates a new ELB frontend
func New(region string, clusterName string, expectedFrontends int) ingress.Frontend {
	log.Infof("ELB Front end region: %s cluster: %s expected frontends: %d", region, clusterName, expectedFrontends)
	metadata := ec2metadata.New(session.New())
	return &elb{
		metadata:          metadata,
		awsElb:            aws_elb.New(session.New(&aws.Config{Region: &region})),
		clusterName:       clusterName,
		region:            region,
		expectedFrontends: expectedFrontends,
		maxTagQuery:       20,
	}
}

type elb struct {
	awsElb              ELB
	metadata            EC2Metadata
	clusterName         string
	region              string
	expectedFrontends   int
	maxTagQuery         int
	instanceID          string
	elbs                []string
	registeredFrontends int
	sync.Mutex
}

// ELB interface to allow mocking of real calls to AWS as well as cutting down the methods from the real
// interface to only the ones we use
type ELB interface {
	DescribeLoadBalancers(input *aws_elb.DescribeLoadBalancersInput) (*aws_elb.DescribeLoadBalancersOutput, error)
	DescribeTags(input *aws_elb.DescribeTagsInput) (*aws_elb.DescribeTagsOutput, error)
	RegisterInstancesWithLoadBalancer(input *aws_elb.RegisterInstancesWithLoadBalancerInput) (*aws_elb.RegisterInstancesWithLoadBalancerOutput, error)
	DeregisterInstancesFromLoadBalancer(input *aws_elb.DeregisterInstancesFromLoadBalancerInput) (*aws_elb.DeregisterInstancesFromLoadBalancerOutput, error)
}

// EC2Metadata interface to allow mocking of the real calls to AWS
type EC2Metadata interface {
	Available() bool
	Region() (string, error)
	GetInstanceIdentityDocument() (ec2metadata.EC2InstanceIdentityDocument, error)
}

func (e *elb) Attach() error {
	e.Lock()
	defer e.Unlock()

	if e.expectedFrontends == 0 {
		return nil
	}

	id, err := e.metadata.GetInstanceIdentityDocument()
	if err != nil {
		return fmt.Errorf("unable to query ec2 metadata service for InstanceId: %v", err)
	}

	instance := id.InstanceID
	log.Infof("Attaching to ELBs from instance %s", instance)
	clusterFrontEnds, err := e.findFrontEndElbs()

	if err != nil {
		return err
	}

	log.Infof("Found %d front ends", len(clusterFrontEnds))

	// Save these now so we can always know what we might have done
	// up until this point we have only read data
	e.elbs = clusterFrontEnds
	e.instanceID = instance
	registered := 0

	for _, frontend := range clusterFrontEnds {
		log.Infof("Registering instance %s with elb %", instance, frontend)
		_, err = e.awsElb.RegisterInstancesWithLoadBalancer(&aws_elb.RegisterInstancesWithLoadBalancerInput{
			Instances: []*aws_elb.Instance{
				&aws_elb.Instance{
					InstanceId: aws.String(instance),
				}},
			LoadBalancerName: aws.String(frontend),
		})

		if err != nil {
			return fmt.Errorf("unable to register instance %s with elb %s: %v", instance, frontend, err)
		}
		registered++

	}

	prometheus.Register(attachedFrontendGauge)
	attachedFrontendGauge.Set(float64(registered))
	e.registeredFrontends = registered
	return nil
}

func (e *elb) findFrontEndElbs() ([]string, error) {
	// Find the load balancers that are tagged with this cluster name
	request := &aws_elb.DescribeLoadBalancersInput{}
	var lbNames []*string
	for {
		resp, err := e.awsElb.DescribeLoadBalancers(request)

		if err != nil {
			return nil, fmt.Errorf("unable to describe load balancers: %v", err)
		}

		for _, entry := range resp.LoadBalancerDescriptions {
			lbNames = append(lbNames, entry.LoadBalancerName)
		}

		if resp.NextMarker == nil {
			break
		}

		// Set the next marker
		request = &aws_elb.DescribeLoadBalancersInput{
			Marker: resp.NextMarker,
		}
	}

	log.Infof("Found %d loadbalancers. Checking for %s tag set to %s", len(lbNames), elbTag, e.clusterName)
	var clusterFrontEnds []string
	totalLbs := len(lbNames)
	for i := 0; i < len(lbNames); i += e.maxTagQuery {
		to := min(i+e.maxTagQuery, totalLbs)
		log.Debugf("Querying tags from %d to %d", i, to)
		names := lbNames[i:to]
		output, err := e.awsElb.DescribeTags(&aws_elb.DescribeTagsInput{
			LoadBalancerNames: names,
		})

		if err != nil {
			return nil, fmt.Errorf("unable to describe tags: %v", err)
		}

		for _, description := range output.TagDescriptions {
			for _, tag := range description.Tags {
				if *tag.Key == elbTag && *tag.Value == e.clusterName {
					log.Infof("Found frontend elb %s", *description.LoadBalancerName)
					clusterFrontEnds = append(clusterFrontEnds, *description.LoadBalancerName)
				}
			}
		}
	}
	return clusterFrontEnds, nil
}

// Detach removes this instance from all the front end ELBs
func (e *elb) Detach() error {
	var failed = false
	for _, elb := range e.elbs {
		log.Infof("Deregistering instance %s with elb %s", e.instanceID, elb)
		_, err := e.awsElb.DeregisterInstancesFromLoadBalancer(&aws_elb.DeregisterInstancesFromLoadBalancerInput{
			Instances:        []*aws_elb.Instance{&aws_elb.Instance{InstanceId: aws.String(e.instanceID)}},
			LoadBalancerName: aws.String(elb),
		})

		if err != nil {
			log.Warnf("unable to deregister instance %s with elb %s: %v", e.instanceID, elb, err)
			failed = true
		}
	}
	if failed {
		return errors.New("at least one ELB failed to detach")
	}
	return nil
}

func (e *elb) Health() error {
	// Lock required to ensure visibility of e.registeredFrontends
	e.Lock()
	defer e.Unlock()
	if e.registeredFrontends != e.expectedFrontends {
		return fmt.Errorf("expected frontends %d registered frontends %d", e.expectedFrontends, e.registeredFrontends)
	}
	return nil
}

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}
<|endoftext|>


<commit_before>// Copyright 2015 Igor Dolzhikov. All rights reserved.
// Use of this source code is governed by
// license that can be found in the LICENSE file.

/*
Package daemon 0.2.10 for use with Go (golang) services.

Package daemon provides primitives for daemonization of golang services.
This package is not provide implementation of user daemon,
accordingly must have root rights to install/remove service.
In the current implementation is only supported Linux and Mac Os X daemon.

Example:

	// Example of a daemon with echo service
	package main

	import (
		"fmt"
		"log"
		"net"
		"os"
		"os/signal"
		"syscall"

		"github.com/takama/daemon"
	)

	const (

		// name of the service, match with executable file name
		name        = "myservice"
		description = "My Echo Service"

		// port which daemon should be listen
		port = ":9977"
	)

	var stdlog, errlog *log.Logger

	// Service has embedded daemon
	type Service struct {
		daemon.Daemon
	}

	// Manage by daemon commands or run the daemon
	func (service *Service) Manage() (string, error) {

		usage := "Usage: myservice install | remove | start | stop | status"

		// if received any kind of command, do it
		if len(os.Args) > 1 {
			command := os.Args[1]
			switch command {
			case "install":
				return service.Install()
			case "remove":
				return service.Remove()
			case "start":
				return service.Start()
			case "stop":
				return service.Stop()
			case "status":
				return service.Status()
			default:
				return usage, nil
			}
		}

		// Do something, call your goroutines, etc

		// Set up channel on which to send signal notifications.
		// We must use a buffered channel or risk missing the signal
		// if we're not ready to receive when the signal is sent.
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

		// Set up listener for defined host and port
		listener, err := net.Listen("tcp", port)
		if err != nil {
			return "Possibly was a problem with the port binding", err
		}

		// set up channel on which to send accepted connections
		listen := make(chan net.Conn, 100)
		go acceptConnection(listener, listen)

		// loop work cycle with accept connections or interrupt
		// by system signal
		for {
			select {
			case conn := <-listen:
				go handleClient(conn)
			case killSignal := <-interrupt:
				stdlog.Println("Got signal:", killSignal)
				stdlog.Println("Stoping listening on ", listener.Addr())
				listener.Close()
				if killSignal == os.Interrupt {
					return "Daemon was interruped by system signal", nil
				}
				return "Daemon was killed", nil
			}
		}

		// never happen, but need to complete code
		return usage, nil
	}

	// Accept a client connection and collect it in a channel
	func acceptConnection(listener net.Listener, listen chan<- net.Conn) {
		for {
			conn, err := listener.Accept()
			if err != nil {
				continue
			}
			listen <- conn
		}
	}

	func handleClient(client net.Conn) {
		for {
			buf := make([]byte, 4096)
			numbytes, err := client.Read(buf)
			if numbytes == 0 || err != nil {
				return
			}
			client.Write(buf)
		}
	}

	func init() {
		stdlog = log.New(os.Stdout, "", log.Ldate|log.Ltime)
		errlog = log.New(os.Stderr, "", log.Ldate|log.Ltime)
	}

	func main() {
		srv, err := daemon.New(name, description)
		if err != nil {
			errlog.Println("Error: ", err)
			os.Exit(1)
		}
		service := &Service{srv}
		status, err := service.Manage()
		if err != nil {
			errlog.Println(status, "\nError: ", err)
			os.Exit(1)
		}
		fmt.Println(status)
	}

Go daemon
*/
package daemon

// Daemon interface has standard set of a methods/commands
type Daemon interface {

	// Install the service into the system
	Install() (string, error)

	// Remove the service and all corresponded files from the system
	Remove() (string, error)

	// Start the service
	Start() (string, error)

	// Stop the service
	Stop() (string, error)

	// Status - check the service status
	Status() (string, error)
}

// New - Create a new daemon
//
// name: name of the service, match with executable file name;
// description: any explanation, what is the service, its purpose
func New(name, description string) (Daemon, error) {
	return newDaemon(name, description)
}
<commit_msg>Bumped version number to 0.2.11<commit_after>// Copyright 2015 Igor Dolzhikov. All rights reserved.
// Use of this source code is governed by
// license that can be found in the LICENSE file.

/*
Package daemon 0.2.11 for use with Go (golang) services.

Package daemon provides primitives for daemonization of golang services.
This package is not provide implementation of user daemon,
accordingly must have root rights to install/remove service.
In the current implementation is only supported Linux and Mac Os X daemon.

Example:

	// Example of a daemon with echo service
	package main

	import (
		"fmt"
		"log"
		"net"
		"os"
		"os/signal"
		"syscall"

		"github.com/takama/daemon"
	)

	const (

		// name of the service, match with executable file name
		name        = "myservice"
		description = "My Echo Service"

		// port which daemon should be listen
		port = ":9977"
	)

	var stdlog, errlog *log.Logger

	// Service has embedded daemon
	type Service struct {
		daemon.Daemon
	}

	// Manage by daemon commands or run the daemon
	func (service *Service) Manage() (string, error) {

		usage := "Usage: myservice install | remove | start | stop | status"

		// if received any kind of command, do it
		if len(os.Args) > 1 {
			command := os.Args[1]
			switch command {
			case "install":
				return service.Install()
			case "remove":
				return service.Remove()
			case "start":
				return service.Start()
			case "stop":
				return service.Stop()
			case "status":
				return service.Status()
			default:
				return usage, nil
			}
		}

		// Do something, call your goroutines, etc

		// Set up channel on which to send signal notifications.
		// We must use a buffered channel or risk missing the signal
		// if we're not ready to receive when the signal is sent.
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

		// Set up listener for defined host and port
		listener, err := net.Listen("tcp", port)
		if err != nil {
			return "Possibly was a problem with the port binding", err
		}

		// set up channel on which to send accepted connections
		listen := make(chan net.Conn, 100)
		go acceptConnection(listener, listen)

		// loop work cycle with accept connections or interrupt
		// by system signal
		for {
			select {
			case conn := <-listen:
				go handleClient(conn)
			case killSignal := <-interrupt:
				stdlog.Println("Got signal:", killSignal)
				stdlog.Println("Stoping listening on ", listener.Addr())
				listener.Close()
				if killSignal == os.Interrupt {
					return "Daemon was interruped by system signal", nil
				}
				return "Daemon was killed", nil
			}
		}

		// never happen, but need to complete code
		return usage, nil
	}

	// Accept a client connection and collect it in a channel
	func acceptConnection(listener net.Listener, listen chan<- net.Conn) {
		for {
			conn, err := listener.Accept()
			if err != nil {
				continue
			}
			listen <- conn
		}
	}

	func handleClient(client net.Conn) {
		for {
			buf := make([]byte, 4096)
			numbytes, err := client.Read(buf)
			if numbytes == 0 || err != nil {
				return
			}
			client.Write(buf)
		}
	}

	func init() {
		stdlog = log.New(os.Stdout, "", log.Ldate|log.Ltime)
		errlog = log.New(os.Stderr, "", log.Ldate|log.Ltime)
	}

	func main() {
		srv, err := daemon.New(name, description)
		if err != nil {
			errlog.Println("Error: ", err)
			os.Exit(1)
		}
		service := &Service{srv}
		status, err := service.Manage()
		if err != nil {
			errlog.Println(status, "\nError: ", err)
			os.Exit(1)
		}
		fmt.Println(status)
	}

Go daemon
*/
package daemon

// Daemon interface has standard set of a methods/commands
type Daemon interface {

	// Install the service into the system
	Install() (string, error)

	// Remove the service and all corresponded files from the system
	Remove() (string, error)

	// Start the service
	Start() (string, error)

	// Stop the service
	Stop() (string, error)

	// Status - check the service status
	Status() (string, error)
}

// New - Create a new daemon
//
// name: name of the service, match with executable file name;
// description: any explanation, what is the service, its purpose
func New(name, description string) (Daemon, error) {
	return newDaemon(name, description)
}
<|endoftext|>