<commit_before>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"time"

	cmds "github.com/jbenet/go-ipfs/commands"
	notif "github.com/jbenet/go-ipfs/notifications"
	peer "github.com/jbenet/go-ipfs/p2p/peer"
	ipdht "github.com/jbenet/go-ipfs/routing/dht"
	u "github.com/jbenet/go-ipfs/util"
)

var ErrNotDHT = errors.New("routing service is not a DHT")

var DhtCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline:          "Issue commands directly through the DHT",
		ShortDescription: ``,
	},

	Subcommands: map[string]*cmds.Command{
		"query":     queryDhtCmd,
		"findprovs": findProvidersDhtCmd,
		"findpeer":  findPeerDhtCmd,
	},
}

var queryDhtCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline:          "Run a 'findClosestPeers' query through the DHT",
		ShortDescription: ``,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("peerID", true, true, "The peerID to run the query against"),
	},
	Options: []cmds.Option{
		cmds.BoolOption("verbose", "v", "Write extra information"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		dht, ok := n.Routing.(*ipdht.IpfsDHT)
		if !ok {
			res.SetError(ErrNotDHT, cmds.ErrNormal)
			return
		}

		events := make(chan *notif.QueryEvent)
		ctx := notif.RegisterForQueryEvents(req.Context().Context, events)

		closestPeers, err := dht.GetClosestPeers(ctx, u.Key(req.Arguments()[0]))

		go func() {
			defer close(events)
			for p := range closestPeers {
				notif.PublishQueryEvent(ctx, &notif.QueryEvent{
					ID:   p,
					Type: notif.FinalPeer,
				})
			}
		}()

		outChan := make(chan interface{})
		res.SetOutput((<-chan interface{})(outChan))

		go func() {
			defer close(outChan)
			for e := range events {
				outChan <- e
			}
		}()
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			marshal := func(v interface{}) (io.Reader, error) {
				obj, ok := v.(*notif.QueryEvent)
				if !ok {
					return nil, u.ErrCast()
				}

				verbose, _, _ := res.Request().Option("v").Bool()

				buf := new(bytes.Buffer)
				if verbose {
					fmt.Fprintf(buf, "%s: ", time.Now().Format("15:04:05.000"))
				}
				switch obj.Type {
				case notif.FinalPeer:
					fmt.Fprintf(buf, "%s\n", obj.ID)
				case notif.PeerResponse:
					if verbose {
						fmt.Fprintf(buf, "* %s says use ", obj.ID)
						for _, p := range obj.Responses {
							fmt.Fprintf(buf, "%s ", p.ID)
						}
						fmt.Fprintln(buf)
					}
				case notif.SendingQuery:
					if verbose {
						fmt.Fprintf(buf, "* querying %s\n", obj.ID)
					}
				case notif.QueryError:
					fmt.Fprintf(buf, "error: %s\n", obj.Extra)
				default:
					fmt.Fprintf(buf, "unrecognized event type: %d\n", obj.Type)
				}
				return buf, nil
			}

			return &cmds.ChannelMarshaler{
				Channel:   outChan,
				Marshaler: marshal,
			}, nil
		},
	},
	Type: notif.QueryEvent{},
}

var findProvidersDhtCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Run a 'FindProviders' query through the DHT",
		ShortDescription: `
FindProviders will return a list of peers who are able to provide the value requested.
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, true, "The key to find providers for"),
	},
	Options: []cmds.Option{
		cmds.BoolOption("verbose", "v", "Write extra information"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		dht, ok := n.Routing.(*ipdht.IpfsDHT)
		if !ok {
			res.SetError(ErrNotDHT, cmds.ErrNormal)
			return
		}

		numProviders := 20

		outChan := make(chan interface{})
		res.SetOutput((<-chan interface{})(outChan))

		events := make(chan *notif.QueryEvent)
		ctx := notif.RegisterForQueryEvents(req.Context().Context, events)

		pchan := dht.FindProvidersAsync(ctx, u.B58KeyDecode(req.Arguments()[0]), numProviders)
		go func() {
			defer close(outChan)
			for e := range events {
				outChan <- e
			}
		}()

		go func() {
			defer close(events)
			for p := range pchan {
				np := p
				notif.PublishQueryEvent(ctx, &notif.QueryEvent{
					Type:      notif.Provider,
					Responses: []*peer.PeerInfo{&np},
				})
			}
		}()
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			verbose, _, _ := res.Request().Option("v").Bool()

			marshal := func(v interface{}) (io.Reader, error) {
				obj, ok := v.(*notif.QueryEvent)
				if !ok {
					return nil, u.ErrCast()
				}

				buf := new(bytes.Buffer)
				if verbose {
					fmt.Fprintf(buf, "%s: ", time.Now().Format("15:04:05.000"))
				}
				switch obj.Type {
				case notif.FinalPeer:
					if verbose {
						fmt.Fprintf(buf, "* closest peer %s\n", obj.ID)
					}
				case notif.Provider:
					fmt.Fprintf(buf, "%s\n", obj.ID.Pretty())
				case notif.PeerResponse:
					if verbose {
						fmt.Fprintf(buf, "* %s says use ", obj.ID)
						for _, p := range obj.Responses {
							fmt.Fprintf(buf, "%s ", p.ID)
						}
						fmt.Fprintln(buf)
					}
				case notif.SendingQuery:
					if verbose {
						fmt.Fprintf(buf, "* querying %s\n", obj.ID)
					}
				case notif.QueryError:
					fmt.Fprintf(buf, "error: %s\n", obj.Extra)
				default:
					fmt.Fprintf(buf, "unrecognized event type: %d\n", obj.Type)
				}
				return buf, nil
			}

			return &cmds.ChannelMarshaler{
				Channel:   outChan,
				Marshaler: marshal,
			}, nil
		},
	},
	Type: notif.QueryEvent{},
}

var findPeerDhtCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline:          "Run a 'FindPeer' query through the DHT",
		ShortDescription: ``,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("peerID", true, true, "The peer to search for"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		dht, ok := n.Routing.(*ipdht.IpfsDHT)
		if !ok {
			res.SetError(ErrNotDHT, cmds.ErrNormal)
			return
		}

		pid, err := peer.IDB58Decode(req.Arguments()[0])
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		outChan := make(chan interface{})
		res.SetOutput((<-chan interface{})(outChan))

		events := make(chan *notif.QueryEvent)
		ctx := notif.RegisterForQueryEvents(req.Context().Context, events)

		go func() {
			defer close(outChan)
			for v := range events {
				outChan <- v
			}
		}()

		go func() {
			defer close(events)
			pi, err := dht.FindPeer(ctx, pid)
			if err != nil {
				notif.PublishQueryEvent(ctx, &notif.QueryEvent{
					Type:  notif.QueryError,
					Extra: err.Error(),
				})
				return
			}

			notif.PublishQueryEvent(ctx, &notif.QueryEvent{
				Type:      notif.FinalPeer,
				Responses: []*peer.PeerInfo{&pi},
			})
		}()
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			marshal := func(v interface{}) (io.Reader, error) {
				obj, ok := v.(*notif.QueryEvent)
				if !ok {
					return nil, u.ErrCast()
				}

				buf := new(bytes.Buffer)
				fmt.Fprintf(buf, "%s: ", time.Now().Format("15:04:05.000"))
				switch obj.Type {
				case notif.FinalPeer:
					pi := obj.Responses[0]
					fmt.Fprintf(buf, "%s\n", pi.ID)
					for _, a := range pi.Addrs {
						fmt.Fprintf(buf, "\t%s\n", a)
					}
				case notif.PeerResponse:
					fmt.Fprintf(buf, "* %s says use ", obj.ID)
					for _, p := range obj.Responses {
						fmt.Fprintf(buf, "%s ", p.ID)
					}
					fmt.Fprintln(buf)
				case notif.SendingQuery:
					fmt.Fprintf(buf, "* querying %s\n", obj.ID)
				case notif.QueryError:
					fmt.Fprintf(buf, "error: %s\n", obj.Extra)
				default:
					fmt.Fprintf(buf, "unrecognized event type: %d\n", obj.Type)
				}
				return buf, nil
			}

			return &cmds.ChannelMarshaler{
				Channel:   outChan,
				Marshaler: marshal,
			}, nil
		},
	},
	Type: notif.QueryEvent{},
}
<commit_msg>print out correct object for dht findproviders command<commit_after>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"time"

	cmds "github.com/jbenet/go-ipfs/commands"
	notif "github.com/jbenet/go-ipfs/notifications"
	peer "github.com/jbenet/go-ipfs/p2p/peer"
	ipdht "github.com/jbenet/go-ipfs/routing/dht"
	u "github.com/jbenet/go-ipfs/util"
)

var ErrNotDHT = errors.New("routing service is not a DHT")

var DhtCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline:          "Issue commands directly through the DHT",
		ShortDescription: ``,
	},

	Subcommands: map[string]*cmds.Command{
		"query":     queryDhtCmd,
		"findprovs": findProvidersDhtCmd,
		"findpeer":  findPeerDhtCmd,
	},
}

var queryDhtCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline:          "Run a 'findClosestPeers' query through the DHT",
		ShortDescription: ``,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("peerID", true, true, "The peerID to run the query against"),
	},
	Options: []cmds.Option{
		cmds.BoolOption("verbose", "v", "Write extra information"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		dht, ok := n.Routing.(*ipdht.IpfsDHT)
		if !ok {
			res.SetError(ErrNotDHT, cmds.ErrNormal)
			return
		}

		events := make(chan *notif.QueryEvent)
		ctx := notif.RegisterForQueryEvents(req.Context().Context, events)

		closestPeers, err := dht.GetClosestPeers(ctx, u.Key(req.Arguments()[0]))

		go func() {
			defer close(events)
			for p := range closestPeers {
				notif.PublishQueryEvent(ctx, &notif.QueryEvent{
					ID:   p,
					Type: notif.FinalPeer,
				})
			}
		}()

		outChan := make(chan interface{})
		res.SetOutput((<-chan interface{})(outChan))

		go func() {
			defer close(outChan)
			for e := range events {
				outChan <- e
			}
		}()
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			marshal := func(v interface{}) (io.Reader, error) {
				obj, ok := v.(*notif.QueryEvent)
				if !ok {
					return nil, u.ErrCast()
				}

				verbose, _, _ := res.Request().Option("v").Bool()

				buf := new(bytes.Buffer)
				if verbose {
					fmt.Fprintf(buf, "%s: ", time.Now().Format("15:04:05.000"))
				}
				switch obj.Type {
				case notif.FinalPeer:
					fmt.Fprintf(buf, "%s\n", obj.ID)
				case notif.PeerResponse:
					if verbose {
						fmt.Fprintf(buf, "* %s says use ", obj.ID)
						for _, p := range obj.Responses {
							fmt.Fprintf(buf, "%s ", p.ID)
						}
						fmt.Fprintln(buf)
					}
				case notif.SendingQuery:
					if verbose {
						fmt.Fprintf(buf, "* querying %s\n", obj.ID)
					}
				case notif.QueryError:
					fmt.Fprintf(buf, "error: %s\n", obj.Extra)
				default:
					fmt.Fprintf(buf, "unrecognized event type: %d\n", obj.Type)
				}
				return buf, nil
			}

			return &cmds.ChannelMarshaler{
				Channel:   outChan,
				Marshaler: marshal,
			}, nil
		},
	},
	Type: notif.QueryEvent{},
}

var findProvidersDhtCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Run a 'FindProviders' query through the DHT",
		ShortDescription: `
FindProviders will return a list of peers who are able to provide the value requested.
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, true, "The key to find providers for"),
	},
	Options: []cmds.Option{
		cmds.BoolOption("verbose", "v", "Write extra information"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		dht, ok := n.Routing.(*ipdht.IpfsDHT)
		if !ok {
			res.SetError(ErrNotDHT, cmds.ErrNormal)
			return
		}

		numProviders := 20

		outChan := make(chan interface{})
		res.SetOutput((<-chan interface{})(outChan))

		events := make(chan *notif.QueryEvent)
		ctx := notif.RegisterForQueryEvents(req.Context().Context, events)

		pchan := dht.FindProvidersAsync(ctx, u.B58KeyDecode(req.Arguments()[0]), numProviders)
		go func() {
			defer close(outChan)
			for e := range events {
				outChan <- e
			}
		}()

		go func() {
			defer close(events)
			for p := range pchan {
				np := p
				notif.PublishQueryEvent(ctx, &notif.QueryEvent{
					Type:      notif.Provider,
					Responses: []*peer.PeerInfo{&np},
				})
			}
		}()
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			verbose, _, _ := res.Request().Option("v").Bool()

			marshal := func(v interface{}) (io.Reader, error) {
				obj, ok := v.(*notif.QueryEvent)
				if !ok {
					return nil, u.ErrCast()
				}

				buf := new(bytes.Buffer)
				if verbose {
					fmt.Fprintf(buf, "%s: ", time.Now().Format("15:04:05.000"))
				}
				switch obj.Type {
				case notif.FinalPeer:
					if verbose {
						fmt.Fprintf(buf, "* closest peer %s\n", obj.ID)
					}
				case notif.Provider:
					prov := obj.Responses[0]
					if verbose {
						fmt.Fprintf(buf, "provider: ")
					}
					fmt.Fprintf(buf, "%s\n", prov.ID.Pretty())
					if verbose {
						for _, a := range prov.Addrs {
							fmt.Fprintf(buf, "\t%s\n", a)
						}
					}
				case notif.PeerResponse:
					if verbose {
						fmt.Fprintf(buf, "* %s says use ", obj.ID)
						for _, p := range obj.Responses {
							fmt.Fprintf(buf, "%s ", p.ID)
						}
						fmt.Fprintln(buf)
					}
				case notif.SendingQuery:
					if verbose {
						fmt.Fprintf(buf, "* querying %s\n", obj.ID)
					}
				case notif.QueryError:
					fmt.Fprintf(buf, "error: %s\n", obj.Extra)
				default:
					fmt.Fprintf(buf, "unrecognized event type: %d\n", obj.Type)
				}
				return buf, nil
			}

			return &cmds.ChannelMarshaler{
				Channel:   outChan,
				Marshaler: marshal,
			}, nil
		},
	},
	Type: notif.QueryEvent{},
}

var findPeerDhtCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline:          "Run a 'FindPeer' query through the DHT",
		ShortDescription: ``,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("peerID", true, true, "The peer to search for"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		dht, ok := n.Routing.(*ipdht.IpfsDHT)
		if !ok {
			res.SetError(ErrNotDHT, cmds.ErrNormal)
			return
		}

		pid, err := peer.IDB58Decode(req.Arguments()[0])
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		outChan := make(chan interface{})
		res.SetOutput((<-chan interface{})(outChan))

		events := make(chan *notif.QueryEvent)
		ctx := notif.RegisterForQueryEvents(req.Context().Context, events)

		go func() {
			defer close(outChan)
			for v := range events {
				outChan <- v
			}
		}()

		go func() {
			defer close(events)
			pi, err := dht.FindPeer(ctx, pid)
			if err != nil {
				notif.PublishQueryEvent(ctx, &notif.QueryEvent{
					Type:  notif.QueryError,
					Extra: err.Error(),
				})
				return
			}

			notif.PublishQueryEvent(ctx, &notif.QueryEvent{
				Type:      notif.FinalPeer,
				Responses: []*peer.PeerInfo{&pi},
			})
		}()
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			marshal := func(v interface{}) (io.Reader, error) {
				obj, ok := v.(*notif.QueryEvent)
				if !ok {
					return nil, u.ErrCast()
				}

				buf := new(bytes.Buffer)
				fmt.Fprintf(buf, "%s: ", time.Now().Format("15:04:05.000"))
				switch obj.Type {
				case notif.FinalPeer:
					pi := obj.Responses[0]
					fmt.Fprintf(buf, "%s\n", pi.ID)
					for _, a := range pi.Addrs {
						fmt.Fprintf(buf, "\t%s\n", a)
					}
				case notif.PeerResponse:
					fmt.Fprintf(buf, "* %s says use ", obj.ID)
					for _, p := range obj.Responses {
						fmt.Fprintf(buf, "%s ", p.ID)
					}
					fmt.Fprintln(buf)
				case notif.SendingQuery:
					fmt.Fprintf(buf, "* querying %s\n", obj.ID)
				case notif.QueryError:
					fmt.Fprintf(buf, "error: %s\n", obj.Extra)
				default:
					fmt.Fprintf(buf, "unrecognized event type: %d\n", obj.Type)
				}
				return buf, nil
			}

			return &cmds.ChannelMarshaler{
				Channel:   outChan,
				Marshaler: marshal,
			}, nil
		},
	},
	Type: notif.QueryEvent{},
}
<|endoftext|>


<commit_before>package core

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
	"time"
	"github.com/cmcoffee/go-snuglib/iotimeout"
)

func (K *KWAPI) Authenticate(username string) (*KWSession, error) {
	return K.authenticate(username, true, false)
}

func (K *KWAPI) AuthLoop(username string) (*KWSession, error) {
	return K.authenticate(username, true, true)
}

// Set User Credentials for kw_api.
func (K *KWAPI) authenticate(username string, permit_change, auth_loop bool) (*KWSession, error) {
	K.testTokenStore()

	if username != NONE {
		session := K.Session(username)
		token, err := K.TokenStore.Load(username)
		if err != nil {
			return nil, err
		}
		if token != nil {
			if token.Expires < time.Now().Add(time.Duration(5*time.Minute)).Unix() {
				// First attempt to use a refresh token if there is one.
				token, err = K.refreshToken(username, token)
				if err != nil {
					if K.secrets.signature_key == nil {
						Notice("Unable to use refresh token, must reauthenticate for new access token: %s", err.Error())
					}
					token = nil
				} else {
					if err := K.TokenStore.Save(username, token); err != nil {
						return &session, err
					}
					return &session, nil
				}
			} else {
				return &session, nil
			}
		}
	}

	if K.secrets.signature_key == nil {
		Stdout("### %s authentication ###\n\n", K.Server)
		for {
			if username == NONE || permit_change {
				username = strings.ToLower(GetInput("-> E-MAIL: "))
			} else {
				                             Stdout("-> E-MAIL: %s", username)
			}
			password :=                   GetSecret("-> PASSWD: ")
			if password == NONE {
				err := fmt.Errorf("Blank password provided.")
				if !auth_loop {
					return nil, err
				}
				Stdout("\n")
				Err("Blank password provided.\n\n")
				continue
			}

			auth, err := K.newToken(username, password)
			if err != nil {
				if !auth_loop {
					return nil, err
				}
				Stdout("\n")
				Err(fmt.Sprintf("%s\n\n", err.Error()))
				continue
			} else {
				session := K.Session(username)
				if err := K.TokenStore.Save(username, auth); err != nil {
					return &session, err
				}
				Stdout("\n")
				return &session, nil
			}
		}
	} else {
		auth, err := K.newToken(username, NONE)
		if err != nil {
			return nil, err
		}
		session := K.Session(username)
		if err := K.TokenStore.Save(username, auth); err != nil {
			return &session, err
		}
		return &session, nil
	}

	return nil, fmt.Errorf("Unable to obtain a token for specified user.")
}

// Add Bearer token to KWAPI requests.
func (s KWSession) setToken(req *http.Request, clear bool) (err error) {
	s.testTokenStore()

	token, err := s.TokenStore.Load(s.Username)
	if err != nil {
		return err
	}

	// If we find a token, check if it's still valid within the next 5 minutes.
	if token != nil && !clear {
		if token.Expires < time.Now().Add(time.Duration(5*time.Minute)).Unix() {
			// First attempt to use a refresh token if there is one.
			token, err = s.refreshToken(s.Username, token)
			if err != nil && s.secrets.signature_key == nil {
				Notice("Unable to use refresh token, must reauthenticate for new access token: %s", err.Error())
			}
		}
	}

	if clear {
		token = nil
		s.TokenStore.Delete(s.Username)
	}

	if token == nil {
		if s.secrets.signature_key != nil {
			token, err = s.newToken(s.Username, NONE)
			if err != nil {
				return err
			}
		} else {
			user, err := s.authenticate(s.Username, false, false)
			if err != nil {
				return err
			}
			token, err = s.TokenStore.Load(user.Username)
			if err != nil {
				return err
			}
		}
	}

	if token != nil {
		req.Header.Set("Authorization", "Bearer "+token.AccessToken)
		if err := s.TokenStore.Save(s.Username, token); err != nil {
			return err
		}
	}
	return nil
}

// Get a new token from a refresh token.
func (K *KWAPI) refreshToken(username string, auth *KWAuth) (*KWAuth, error) {
	if auth == nil {
		return nil, fmt.Errorf("No refresh token found for %s.", username)
	}
	path := fmt.Sprintf("https://%s/oauth/token", K.Server)

	req, err := http.NewRequest(http.MethodPost, path, nil)
	if err != nil {
		return nil, err
	}

	http_header := make(http.Header)
	http_header.Set("Content-Type", "application/x-www-form-urlencoded")
	if K.AgentString == NONE {
		K.AgentString = "SnugLib/1.0"
	}
	http_header.Set("User-Agent", K.AgentString)

	req.Header = http_header

	client_id := K.ApplicationID

	postform := &url.Values{
		"client_id":     {client_id},
		"client_secret": {K.secrets.decrypt(K.secrets.client_secret_key)},
		"grant_type":    {"refresh_token"},
		"refresh_token": {auth.RefreshToken},
	}

	if K.Snoop {
		Debug("[kiteworks]: %s", username)
		Debug("--> ACTION: \"POST\" PATH: \"%s\"", path)
		for k, v := range *postform {
			if k == "grant_type" || k == "RedirectURI" || k == "scope" {
				Debug("\\-> POST PARAM: %s VALUE: %s", k, v)
			} else {
				Debug("\\-> POST PARAM: %s VALUE: [HIDDEN]", k)
			}
		}
	}

	req.Body = ioutil.NopCloser(bytes.NewReader([]byte(postform.Encode())))
	req.Body = iotimeout.NewReadCloser(req.Body, K.RequestTimeout)
	defer req.Body.Close()

	client := K.Session(username).NewClient()

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if err := K.decodeJSON(resp, &auth); err != nil {
		return nil, err
	}

	auth.Expires = auth.Expires + time.Now().Unix()
	return auth, nil
}

// Generate a new Bearer token from kiteworks.
func (K *KWAPI) newToken(username, password string) (auth *KWAuth, err error) {

	path := fmt.Sprintf("https://%s/oauth/token", K.Server)

	req, err := http.NewRequest(http.MethodPost, path, nil)
	if err != nil {
		return nil, err
	}

	http_header := make(http.Header)
	http_header.Set("Content-Type", "application/x-www-form-urlencoded")
	http_header.Set("User-Agent", K.AgentString)

	req.Header = http_header

	client_id := K.ApplicationID

	postform := &url.Values{
		"client_id":     {client_id},
		"client_secret": {K.secrets.decrypt(K.secrets.client_secret_key)},
		"redirect_uri":  {K.RedirectURI},
	}

	if password != NONE {
		postform.Add("grant_type", "password")
		postform.Add("username", username)
		postform.Add("password", password)
	} else {
		signature := K.secrets.decrypt(K.secrets.signature_key)
		randomizer := rand.New(rand.NewSource(int64(time.Now().Unix())))
		nonce := randomizer.Int() % 999999
		timestamp := int64(time.Now().Unix())

		base_string := fmt.Sprintf("%s|@@|%s|@@|%d|@@|%d", client_id, username, timestamp, nonce)

		mac := hmac.New(sha1.New, []byte(signature))
		mac.Write([]byte(base_string))
		signature = hex.EncodeToString(mac.Sum(nil))

		auth_code := fmt.Sprintf("%s|@@|%s|@@|%d|@@|%d|@@|%s",
			base64.StdEncoding.EncodeToString([]byte(client_id)),
			base64.StdEncoding.EncodeToString([]byte(username)),
			timestamp, nonce, signature)

		postform.Add("grant_type", "authorization_code")
		postform.Add("code", auth_code)

	}

	if K.Snoop {
		Debug("[kiteworks]: %s", username)
		Debug("--> ACTION: \"POST\" PATH: \"%s\"", path)
		for k, v := range *postform {
			if k == "grant_type" || k == "redirect_uri" || k == "scope" {
				Debug("\\-> POST PARAM: %s VALUE: %s", k, v)
			} else {
				Debug("\\-> POST PARAM: %s VALUE: [HIDDEN]", k)
			}
		}
	}

	req.Body = ioutil.NopCloser(bytes.NewReader([]byte(postform.Encode())))
	req.Body = iotimeout.NewReadCloser(req.Body, K.RequestTimeout)
	defer req.Body.Close()

	client := K.Session(username).NewClient()

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if err := K.decodeJSON(resp, &auth); err != nil {
		return nil, err
	}

	auth.Expires = auth.Expires + time.Now().Unix()
	return
}
<commit_msg>Updated Login Prompt<commit_after>package core

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
	"time"
	"github.com/cmcoffee/go-snuglib/iotimeout"
)

func (K *KWAPI) Authenticate(username string) (*KWSession, error) {
	return K.authenticate(username, true, false)
}

func (K *KWAPI) AuthLoop(username string) (*KWSession, error) {
	return K.authenticate(username, true, true)
}

// Set User Credentials for kw_api.
func (K *KWAPI) authenticate(username string, permit_change, auth_loop bool) (*KWSession, error) {
	K.testTokenStore()

	if username != NONE {
		session := K.Session(username)
		token, err := K.TokenStore.Load(username)
		if err != nil {
			return nil, err
		}
		if token != nil {
			if token.Expires < time.Now().Add(time.Duration(5*time.Minute)).Unix() {
				// First attempt to use a refresh token if there is one.
				token, err = K.refreshToken(username, token)
				if err != nil {
					if K.secrets.signature_key == nil {
						Notice("Unable to use refresh token, must reauthenticate for new access token: %s", err.Error())
					}
					token = nil
				} else {
					if err := K.TokenStore.Save(username, token); err != nil {
						return &session, err
					}
					return &session, nil
				}
			} else {
				return &session, nil
			}
		}
	}

	if K.secrets.signature_key == nil {
		defer PleaseWait.Hide()
		Stdout("### %s authentication ###\n\n", K.Server)
		for {
			PleaseWait.Hide()
			if username == NONE || permit_change {
				username = strings.ToLower(GetInput(" -> Account Login(email): "))
			} else {
				                             Stdout(" -> Account Login(email): %s", username)
			}
			password :=                   GetSecret(" -> Account Password: ")
			if password == NONE {
				err := fmt.Errorf("Blank password provided.")
				if !auth_loop {
					return nil, err
				}
				Stdout("\n")
				Err("Blank password provided.\n\n")
				continue
			}
			PleaseWait.Show()
			auth, err := K.newToken(username, password)
			if err != nil {
				if !auth_loop {
					return nil, err
				}
				Stdout("\n")
				Err(fmt.Sprintf("%s\n\n", err.Error()))
				continue
			} else {
				session := K.Session(username)
				if err := K.TokenStore.Save(username, auth); err != nil {
					return &session, err
				}
				Stdout("\n")
				return &session, nil
			}
		}
	} else {
		auth, err := K.newToken(username, NONE)
		if err != nil {
			return nil, err
		}
		session := K.Session(username)
		if err := K.TokenStore.Save(username, auth); err != nil {
			return &session, err
		}
		return &session, nil
	}

	return nil, fmt.Errorf("Unable to obtain a token for specified user.")
}

// Add Bearer token to KWAPI requests.
func (s KWSession) setToken(req *http.Request, clear bool) (err error) {
	s.testTokenStore()

	token, err := s.TokenStore.Load(s.Username)
	if err != nil {
		return err
	}

	// If we find a token, check if it's still valid within the next 5 minutes.
	if token != nil && !clear {
		if token.Expires < time.Now().Add(time.Duration(5*time.Minute)).Unix() {
			// First attempt to use a refresh token if there is one.
			token, err = s.refreshToken(s.Username, token)
			if err != nil && s.secrets.signature_key == nil {
				Notice("Unable to use refresh token, must reauthenticate for new access token: %s", err.Error())
			}
		}
	}

	if clear {
		token = nil
		s.TokenStore.Delete(s.Username)
	}

	if token == nil {
		if s.secrets.signature_key != nil {
			token, err = s.newToken(s.Username, NONE)
			if err != nil {
				return err
			}
		} else {
			user, err := s.authenticate(s.Username, false, false)
			if err != nil {
				return err
			}
			token, err = s.TokenStore.Load(user.Username)
			if err != nil {
				return err
			}
		}
	}

	if token != nil {
		req.Header.Set("Authorization", "Bearer "+token.AccessToken)
		if err := s.TokenStore.Save(s.Username, token); err != nil {
			return err
		}
	}
	return nil
}

// Get a new token from a refresh token.
func (K *KWAPI) refreshToken(username string, auth *KWAuth) (*KWAuth, error) {
	if auth == nil {
		return nil, fmt.Errorf("No refresh token found for %s.", username)
	}
	path := fmt.Sprintf("https://%s/oauth/token", K.Server)

	req, err := http.NewRequest(http.MethodPost, path, nil)
	if err != nil {
		return nil, err
	}

	http_header := make(http.Header)
	http_header.Set("Content-Type", "application/x-www-form-urlencoded")
	if K.AgentString == NONE {
		K.AgentString = "SnugLib/1.0"
	}
	http_header.Set("User-Agent", K.AgentString)

	req.Header = http_header

	client_id := K.ApplicationID

	postform := &url.Values{
		"client_id":     {client_id},
		"client_secret": {K.secrets.decrypt(K.secrets.client_secret_key)},
		"grant_type":    {"refresh_token"},
		"refresh_token": {auth.RefreshToken},
	}

	if K.Snoop {
		Debug("[kiteworks]: %s", username)
		Debug("--> ACTION: \"POST\" PATH: \"%s\"", path)
		for k, v := range *postform {
			if k == "grant_type" || k == "RedirectURI" || k == "scope" {
				Debug("\\-> POST PARAM: %s VALUE: %s", k, v)
			} else {
				Debug("\\-> POST PARAM: %s VALUE: [HIDDEN]", k)
			}
		}
	}

	req.Body = ioutil.NopCloser(bytes.NewReader([]byte(postform.Encode())))
	req.Body = iotimeout.NewReadCloser(req.Body, K.RequestTimeout)
	defer req.Body.Close()

	client := K.Session(username).NewClient()

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if err := K.decodeJSON(resp, &auth); err != nil {
		return nil, err
	}

	auth.Expires = auth.Expires + time.Now().Unix()
	return auth, nil
}

// Generate a new Bearer token from kiteworks.
func (K *KWAPI) newToken(username, password string) (auth *KWAuth, err error) {

	path := fmt.Sprintf("https://%s/oauth/token", K.Server)

	req, err := http.NewRequest(http.MethodPost, path, nil)
	if err != nil {
		return nil, err
	}

	http_header := make(http.Header)
	http_header.Set("Content-Type", "application/x-www-form-urlencoded")
	http_header.Set("User-Agent", K.AgentString)

	req.Header = http_header

	client_id := K.ApplicationID

	postform := &url.Values{
		"client_id":     {client_id},
		"client_secret": {K.secrets.decrypt(K.secrets.client_secret_key)},
		"redirect_uri":  {K.RedirectURI},
	}

	if password != NONE {
		postform.Add("grant_type", "password")
		postform.Add("username", username)
		postform.Add("password", password)
	} else {
		signature := K.secrets.decrypt(K.secrets.signature_key)
		randomizer := rand.New(rand.NewSource(int64(time.Now().Unix())))
		nonce := randomizer.Int() % 999999
		timestamp := int64(time.Now().Unix())

		base_string := fmt.Sprintf("%s|@@|%s|@@|%d|@@|%d", client_id, username, timestamp, nonce)

		mac := hmac.New(sha1.New, []byte(signature))
		mac.Write([]byte(base_string))
		signature = hex.EncodeToString(mac.Sum(nil))

		auth_code := fmt.Sprintf("%s|@@|%s|@@|%d|@@|%d|@@|%s",
			base64.StdEncoding.EncodeToString([]byte(client_id)),
			base64.StdEncoding.EncodeToString([]byte(username)),
			timestamp, nonce, signature)

		postform.Add("grant_type", "authorization_code")
		postform.Add("code", auth_code)

	}

	if K.Snoop {
		Debug("[kiteworks]: %s", username)
		Debug("--> ACTION: \"POST\" PATH: \"%s\"", path)
		for k, v := range *postform {
			if k == "grant_type" || k == "redirect_uri" || k == "scope" {
				Debug("\\-> POST PARAM: %s VALUE: %s", k, v)
			} else {
				Debug("\\-> POST PARAM: %s VALUE: [HIDDEN]", k)
			}
		}
	}

	req.Body = ioutil.NopCloser(bytes.NewReader([]byte(postform.Encode())))
	req.Body = iotimeout.NewReadCloser(req.Body, K.RequestTimeout)
	defer req.Body.Close()

	client := K.Session(username).NewClient()

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if err := K.decodeJSON(resp, &auth); err != nil {
		return nil, err
	}

	auth.Expires = auth.Expires + time.Now().Unix()
	return
}
<|endoftext|>


<commit_before>package state

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/txn"
)

// Life represents the lifecycle state of the entities
// Relation, Unit, Service and Machine.
type Life int8

const (
	Alive Life = iota
	Dying
	Dead
	nLife
)

var notDead = D{{"life", D{{"$ne", Dead}}}}
var isAlive = D{{"life", Alive}}

var lifeStrings = [nLife]string{
	Alive: "alive",
	Dying: "dying",
	Dead:  "dead",
}

func (l Life) String() string {
	return lifeStrings[l]
}

// Living describes state entities with a lifecycle.
type Living interface {
	Life() Life
	EnsureDying() error
	EnsureDead() error
	Refresh() error
}

// ensureDying advances the specified entity's life status to Dying, if necessary.
func ensureDying(st *State, coll *mgo.Collection, id interface{}, desc string) error {
	ops := []txn.Op{{
		C:      coll.Name,
		Id:     id,
		Assert: isAlive,
		Update: D{{"$set", D{{"life", Dying}}}},
	}}
	if err := st.runner.Run(ops, "", nil); err == txn.ErrAborted {
		return nil
	} else if err != nil {
		return fmt.Errorf("cannot start termination of %s %#v: %v", desc, id, err)
	}
	return nil
}

// cannotKillError is returned from ensureDead when the targeted entity's
// lifecycle has failed to advance to (or beyond) Dead, due to assertion
// failures.
type cannotKillError struct {
	prefix, msg string
}

func (e *cannotKillError) Error() string {
	return fmt.Sprintf("%s: %s", e.prefix, e.msg)
}

// ensureDead advances the specified entity's life status to Dead, if necessary.
// Preconditions can be supplied in assertOps; if the preconditions fail, the error
// will contain assertMsg. If the entity is not found, no error is returned.
func ensureDead(st *State, coll *mgo.Collection, id interface{}, desc string, assertOps []txn.Op, assertMsg string) (err error) {
	errPrefix := fmt.Sprintf("cannot finish termination of %s %#v", desc, id)
	decorate := func(err error) error {
		return fmt.Errorf("%s: %v", errPrefix, err)
	}
	ops := append(assertOps, txn.Op{
		C:      coll.Name,
		Id:     id,
		Update: D{{"$set", D{{"life", Dead}}}},
	})
	if err = st.runner.Run(ops, "", nil); err == nil {
		return nil
	} else if err != txn.ErrAborted {
		return decorate(err)
	}
	var doc struct{ Life }
	if err = coll.FindId(id).One(&doc); err == mgo.ErrNotFound {
		return nil
	} else if err != nil {
		return decorate(err)
	} else if doc.Life != Dead {
		return &cannotKillError{errPrefix, assertMsg}
	}
	return nil
}
<commit_msg>revert life.go, changes weren't good<commit_after>package state

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/trivial"
)

// Life represents the lifecycle state of the entities
// Relation, Unit, Service and Machine.
type Life int8

const (
	Alive Life = iota
	Dying
	Dead
	nLife
)

var notDead = D{{"life", D{{"$ne", Dead}}}}
var isAlive = D{{"life", Alive}}

var lifeStrings = [nLife]string{
	Alive: "alive",
	Dying: "dying",
	Dead:  "dead",
}

func (l Life) String() string {
	return lifeStrings[l]
}

// Living describes state entities with a lifecycle.
type Living interface {
	Life() Life
	EnsureDying() error
	EnsureDead() error
	Refresh() error
}

// ensureDying advances the specified entity's life status to Dying, if necessary.
func ensureDying(st *State, coll *mgo.Collection, id interface{}, desc string) error {
	ops := []txn.Op{{
		C:      coll.Name,
		Id:     id,
		Assert: isAlive,
		Update: D{{"$set", D{{"life", Dying}}}},
	}}
	if err := st.runner.Run(ops, "", nil); err == txn.ErrAborted {
		return nil
	} else if err != nil {
		return fmt.Errorf("cannot start termination of %s %#v: %v", desc, id, err)
	}
	return nil
}

// ensureDead advances the specified entity's life status to Dead, if necessary.
// Preconditions can be supplied in assertOps; if the preconditions fail, the error
// will contain assertMsg. If the entity is not found, no error is returned.
func ensureDead(st *State, coll *mgo.Collection, id interface{}, desc string, assertOps []txn.Op, assertMsg string) (err error) {
	defer trivial.ErrorContextf(&err, "cannot finish termination of %s %#v", desc, id)
	ops := append(assertOps, txn.Op{
		C:      coll.Name,
		Id:     id,
		Update: D{{"$set", D{{"life", Dead}}}},
	})
	if err = st.runner.Run(ops, "", nil); err == nil {
		return nil
	} else if err != txn.ErrAborted {
		return err
	}
	var doc struct{ Life }
	if err = coll.FindId(id).One(&doc); err == mgo.ErrNotFound {
		return nil
	} else if err != nil {
		return err
	} else if doc.Life != Dead {
		return fmt.Errorf(assertMsg)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"github.com/tsuru/tsuru/cmd"
	"launchpad.net/gocheck"
	"os"
	"os/exec"
	"testing"
)

type S struct {
	recover []string
}

func (s *S) SetUpSuite(c *gocheck.C) {
	targetFile := os.Getenv("HOME") + "/.tsuru_target"
	_, err := os.Stat(targetFile)
	if err == nil {
		old := targetFile + ".old"
		s.recover = []string{"mv", old, targetFile}
		exec.Command("mv", targetFile, old).Run()
	} else {
		s.recover = []string{"rm", targetFile}
	}
	f, err := os.Create(targetFile)
	c.Assert(err, gocheck.IsNil)
	f.Write([]byte("http://localhost"))
	f.Close()
}

func (s *S) TearDownSuite(c *gocheck.C) {
	exec.Command(s.recover[0], s.recover[1:]...).Run()
}

var _ = gocheck.Suite(&S{})
var manager *cmd.Manager

func Test(t *testing.T) { gocheck.TestingT(t) }

func (s *S) SetUpTest(c *gocheck.C) {
	var stdout, stderr bytes.Buffer
	manager = cmd.NewManager("glb", version, header, &stdout, &stderr, os.Stdin, nil)
}
<commit_msg>testing: Extracted code to ensure .tsuru_target file exists<commit_after>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"github.com/tsuru/tsuru/cmd"
	tTesting "github.com/tsuru/tsuru/testing"
	"launchpad.net/gocheck"
	"os"
	"testing"
)

type S struct {
	recover []string
}

func (s *S) SetUpSuite(c *gocheck.C) {
	s.recover = tTesting.SetTargetFile(c)
}

func (s *S) TearDownSuite(c *gocheck.C) {
	tTesting.RollbackTargetFile(s.recover)
}

var _ = gocheck.Suite(&S{})
var manager *cmd.Manager

func Test(t *testing.T) { gocheck.TestingT(t) }

func (s *S) SetUpTest(c *gocheck.C) {
	var stdout, stderr bytes.Buffer
	manager = cmd.NewManager("glb", version, header, &stdout, &stderr, os.Stdin, nil)
}
<|endoftext|>


<commit_before>/*
Package supervisor provides supervisor trees for Go applications.

This package is a clean reimplementation of github.com/thejerf/suture, aiming
to be more Go idiomatic, thus less Erlang-like.

It is built on top of context package, with all of its advantages, namely the
possibility trickle down context-related values and cancellation signals.

TheJerf's blog post about Suture is a very good and helpful read to understand
how this package has been implemented.

http://www.jerf.org/iri/post/2930
*/
package supervisor // import "cirello.io/supervisor"

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

	"golang.org/x/net/context"
)

// Service is the public interface expected by a Supervisor.
//
// This will be internally named after the result of fmt.Stringer, if available.
// Otherwise it will going to use an internal representation for the service
// name.
type Service interface {
	// Serve is called by a Supervisor to start the service. It expects the
	// service to honor the passed context and its lifetime. Observe
	// <-ctx.Done() and ctx.Err(). If the service is stopped by anything
	// but the Supervisor, it will get started again. Be careful with shared
	// state among restarts.
	Serve(ctx context.Context)
}

// Supervisor is the basic datastructure responsible for offering a supervisor
// tree. It implements Service, therefore it can be nested if necessary. When
// passing the Supervisor around, remind to do it as reference (&supervisor).
type Supervisor struct {
	// Name for this supervisor tree, used for logging.
	Name string

	// FailureDecay is the timespan on which the current failure count will
	// be halved.
	FailureDecay float64

	// FailureThreshold is the maximum accepted number of failures, after
	// decay adjustment, that shall trigger the back-off wait.
	FailureThreshold float64

	// Backoff is the wait duration when hit threshold.
	Backoff time.Duration

	// Log is a replaceable function used for overall logging
	Log func(string)

	ready sync.Once

	startedMu sync.Mutex
	started   bool

	addedService    chan struct{}
	startedServices chan struct{}
	stoppedService  chan struct{}

	servicesMu sync.Mutex
	services   map[string]Service

	cancellationsMu sync.Mutex
	cancellations   map[string]context.CancelFunc

	backoffMu sync.Mutex
	backoff   map[string]*backoff

	runningMu sync.Mutex
	running   int
}

func (s *Supervisor) String() string {
	return s.Name
}

func (s *Supervisor) prepare() {
	s.ready.Do(func() {
		if s.Name == "" {
			s.Name = "supervisor"
		}
		s.addedService = make(chan struct{}, 1)
		s.backoff = make(map[string]*backoff)
		s.cancellations = make(map[string]context.CancelFunc)
		s.services = make(map[string]Service)
		s.startedServices = make(chan struct{})
		s.stoppedService = make(chan struct{}, 1)

		if s.Log == nil {
			s.Log = func(str string) {
				log.Println(s.Name, ":", str)
			}
		}
		if s.FailureDecay == 0 {
			s.FailureDecay = 30
		}
		if s.FailureThreshold == 0 {
			s.FailureThreshold = 5
		}
		if s.Backoff == 0 {
			s.Backoff = 15 * time.Second
		}
	})
}

// Add inserts into the Supervisor tree a new service. If the Supervisor is
// already started, it will start it automatically.
func (s *Supervisor) Add(service Service) {
	s.prepare()

	name := fmt.Sprintf("%s", service)

	s.servicesMu.Lock()
	s.backoffMu.Lock()
	s.backoff[name] = &backoff{}
	s.services[name] = service
	s.backoffMu.Unlock()
	s.servicesMu.Unlock()

	select {
	case s.addedService <- struct{}{}:
	default:
	}
}

// Remove stops the service in the Supervisor tree and remove from it.
func (s *Supervisor) Remove(name string) {
	s.prepare()

	s.servicesMu.Lock()
	defer s.servicesMu.Unlock()
	if _, ok := s.services[name]; !ok {
		return
	}

	s.cancellationsMu.Lock()
	defer s.cancellationsMu.Unlock()
	if c, ok := s.cancellations[name]; ok {
		delete(s.cancellations, name)
		c()
	}
}

// Serve starts the Supervisor tree. It can be started only once at a time. If
// stopped (canceled), it can be restarted.
func (s *Supervisor) Serve(ctx context.Context) {
	s.prepare()

	select {
	case s.addedService <- struct{}{}:
	default:
	}

	s.startedMu.Lock()
	if !s.started {
		s.started = true
		s.startedMu.Unlock()

		s.serve(ctx)

		s.startedMu.Lock()
		s.started = false
	}
	s.startedMu.Unlock()
}

// Services return a list of services
func (s *Supervisor) Services() map[string]Service {
	svclist := make(map[string]Service)
	s.servicesMu.Lock()
	for k, v := range s.services {
		svclist[k] = v
	}
	s.servicesMu.Unlock()
	return svclist
}

// Cancelations return a list of services names and their cancellation calls
func (s *Supervisor) Cancelations() map[string]context.CancelFunc {
	svclist := make(map[string]context.CancelFunc)
	s.cancellationsMu.Lock()
	for k, v := range s.cancellations {
		svclist[k] = v
	}
	s.cancellationsMu.Unlock()
	return svclist
}

func (s *Supervisor) serve(ctx context.Context) {
	go func(ctx context.Context) {
		for {
			select {
			case <-s.addedService:
				s.startServices(ctx)
				select {
				case s.startedServices <- struct{}{}:
				default:
				}

			case <-ctx.Done():
				return
			}
		}
	}(ctx)

	<-ctx.Done()
	s.cancellationsMu.Lock()
	s.cancellations = make(map[string]context.CancelFunc)
	s.cancellationsMu.Unlock()

	for range s.stoppedService {
		s.runningMu.Lock()
		r := s.running
		s.runningMu.Unlock()
		if r == 0 {
			return
		}
	}
}

func (s *Supervisor) startServices(ctx context.Context) {
	s.servicesMu.Lock()
	defer s.servicesMu.Unlock()

	var wg sync.WaitGroup

	for name, svc := range s.services {
		s.cancellationsMu.Lock()
		_, ok := s.cancellations[name]
		s.cancellationsMu.Unlock()
		if ok {
			continue
		}

		wg.Add(1)

		go func(name string, svc Service) {
			s.runningMu.Lock()
			s.running++
			s.runningMu.Unlock()
			wg.Done()
			for {
				retry := func() (retry bool) {
					select {
					case <-ctx.Done():
						return false
					default:
					}

					defer func() {
						if r := recover(); r != nil {
							s.Log(fmt.Sprint("trapped panic:", r))
							retry = true
						}
					}()

					c, cancel := context.WithCancel(ctx)
					s.cancellationsMu.Lock()
					s.cancellations[name] = cancel
					s.cancellationsMu.Unlock()
					svc.Serve(c)

					select {
					case <-ctx.Done():
						return false
					default:
						return true
					}
				}()
				if retry {
					s.Log(fmt.Sprintf("restarting %s", name))
					s.backoffMu.Lock()
					b := s.backoff[name]
					s.backoffMu.Unlock()
					b.wait(s.FailureDecay, s.FailureThreshold, s.Backoff)
					continue
				}

				s.runningMu.Lock()
				s.running--
				s.runningMu.Unlock()
				s.stoppedService <- struct{}{}
				return
			}
		}(name, svc)
	}
	wg.Wait()
}

type backoff struct {
	lastfail time.Time
	failures float64
}

func (b *backoff) wait(failureDecay float64, threshold float64, backoffDur time.Duration) {
	if b.lastfail.IsZero() {
		b.lastfail = time.Now()
		b.failures = 1.0
	} else {
		b.failures++
		intervals := time.Since(b.lastfail).Seconds() / failureDecay
		b.failures = b.failures*math.Pow(.5, intervals) + 1
	}

	if b.failures > threshold {
		time.Sleep(backoffDur)
	}
}
<commit_msg>cancellation -> cancelation<commit_after>/*
Package supervisor provides supervisor trees for Go applications.

This package is a clean reimplementation of github.com/thejerf/suture, aiming
to be more Go idiomatic, thus less Erlang-like.

It is built on top of context package, with all of its advantages, namely the
possibility trickle down context-related values and cancellation signals.

TheJerf's blog post about Suture is a very good and helpful read to understand
how this package has been implemented.

http://www.jerf.org/iri/post/2930
*/
package supervisor // import "cirello.io/supervisor"

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

	"golang.org/x/net/context"
)

// Service is the public interface expected by a Supervisor.
//
// This will be internally named after the result of fmt.Stringer, if available.
// Otherwise it will going to use an internal representation for the service
// name.
type Service interface {
	// Serve is called by a Supervisor to start the service. It expects the
	// service to honor the passed context and its lifetime. Observe
	// <-ctx.Done() and ctx.Err(). If the service is stopped by anything
	// but the Supervisor, it will get started again. Be careful with shared
	// state among restarts.
	Serve(ctx context.Context)
}

// Supervisor is the basic datastructure responsible for offering a supervisor
// tree. It implements Service, therefore it can be nested if necessary. When
// passing the Supervisor around, remind to do it as reference (&supervisor).
type Supervisor struct {
	// Name for this supervisor tree, used for logging.
	Name string

	// FailureDecay is the timespan on which the current failure count will
	// be halved.
	FailureDecay float64

	// FailureThreshold is the maximum accepted number of failures, after
	// decay adjustment, that shall trigger the back-off wait.
	FailureThreshold float64

	// Backoff is the wait duration when hit threshold.
	Backoff time.Duration

	// Log is a replaceable function used for overall logging
	Log func(string)

	ready sync.Once

	startedMu sync.Mutex
	started   bool

	addedService    chan struct{}
	startedServices chan struct{}
	stoppedService  chan struct{}

	servicesMu sync.Mutex
	services   map[string]Service

	cancelationsMu sync.Mutex
	cancelations   map[string]context.CancelFunc

	backoffMu sync.Mutex
	backoff   map[string]*backoff

	runningMu sync.Mutex
	running   int
}

func (s *Supervisor) String() string {
	return s.Name
}

func (s *Supervisor) prepare() {
	s.ready.Do(func() {
		if s.Name == "" {
			s.Name = "supervisor"
		}
		s.addedService = make(chan struct{}, 1)
		s.backoff = make(map[string]*backoff)
		s.cancelations = make(map[string]context.CancelFunc)
		s.services = make(map[string]Service)
		s.startedServices = make(chan struct{})
		s.stoppedService = make(chan struct{}, 1)

		if s.Log == nil {
			s.Log = func(str string) {
				log.Println(s.Name, ":", str)
			}
		}
		if s.FailureDecay == 0 {
			s.FailureDecay = 30
		}
		if s.FailureThreshold == 0 {
			s.FailureThreshold = 5
		}
		if s.Backoff == 0 {
			s.Backoff = 15 * time.Second
		}
	})
}

// Add inserts into the Supervisor tree a new service. If the Supervisor is
// already started, it will start it automatically.
func (s *Supervisor) Add(service Service) {
	s.prepare()

	name := fmt.Sprintf("%s", service)

	s.servicesMu.Lock()
	s.backoffMu.Lock()
	s.backoff[name] = &backoff{}
	s.services[name] = service
	s.backoffMu.Unlock()
	s.servicesMu.Unlock()

	select {
	case s.addedService <- struct{}{}:
	default:
	}
}

// Remove stops the service in the Supervisor tree and remove from it.
func (s *Supervisor) Remove(name string) {
	s.prepare()

	s.servicesMu.Lock()
	defer s.servicesMu.Unlock()
	if _, ok := s.services[name]; !ok {
		return
	}

	s.cancelationsMu.Lock()
	defer s.cancelationsMu.Unlock()
	if c, ok := s.cancelations[name]; ok {
		delete(s.cancelations, name)
		c()
	}
}

// Serve starts the Supervisor tree. It can be started only once at a time. If
// stopped (canceled), it can be restarted.
func (s *Supervisor) Serve(ctx context.Context) {
	s.prepare()

	select {
	case s.addedService <- struct{}{}:
	default:
	}

	s.startedMu.Lock()
	if !s.started {
		s.started = true
		s.startedMu.Unlock()

		s.serve(ctx)

		s.startedMu.Lock()
		s.started = false
	}
	s.startedMu.Unlock()
}

// Services return a list of services
func (s *Supervisor) Services() map[string]Service {
	svclist := make(map[string]Service)
	s.servicesMu.Lock()
	for k, v := range s.services {
		svclist[k] = v
	}
	s.servicesMu.Unlock()
	return svclist
}

// Cancelations return a list of services names and their cancellation calls
func (s *Supervisor) Cancelations() map[string]context.CancelFunc {
	svclist := make(map[string]context.CancelFunc)
	s.cancelationsMu.Lock()
	for k, v := range s.cancelations {
		svclist[k] = v
	}
	s.cancelationsMu.Unlock()
	return svclist
}

func (s *Supervisor) serve(ctx context.Context) {
	go func(ctx context.Context) {
		for {
			select {
			case <-s.addedService:
				s.startServices(ctx)
				select {
				case s.startedServices <- struct{}{}:
				default:
				}

			case <-ctx.Done():
				return
			}
		}
	}(ctx)

	<-ctx.Done()
	s.cancelationsMu.Lock()
	s.cancelations = make(map[string]context.CancelFunc)
	s.cancelationsMu.Unlock()

	for range s.stoppedService {
		s.runningMu.Lock()
		r := s.running
		s.runningMu.Unlock()
		if r == 0 {
			return
		}
	}
}

func (s *Supervisor) startServices(ctx context.Context) {
	s.servicesMu.Lock()
	defer s.servicesMu.Unlock()

	var wg sync.WaitGroup

	for name, svc := range s.services {
		s.cancelationsMu.Lock()
		_, ok := s.cancelations[name]
		s.cancelationsMu.Unlock()
		if ok {
			continue
		}

		wg.Add(1)

		go func(name string, svc Service) {
			s.runningMu.Lock()
			s.running++
			s.runningMu.Unlock()
			wg.Done()
			for {
				retry := func() (retry bool) {
					select {
					case <-ctx.Done():
						return false
					default:
					}

					defer func() {
						if r := recover(); r != nil {
							s.Log(fmt.Sprint("trapped panic:", r))
							retry = true
						}
					}()

					c, cancel := context.WithCancel(ctx)
					s.cancelationsMu.Lock()
					s.cancelations[name] = cancel
					s.cancelationsMu.Unlock()
					svc.Serve(c)

					select {
					case <-ctx.Done():
						return false
					default:
						return true
					}
				}()
				if retry {
					s.Log(fmt.Sprintf("restarting %s", name))
					s.backoffMu.Lock()
					b := s.backoff[name]
					s.backoffMu.Unlock()
					b.wait(s.FailureDecay, s.FailureThreshold, s.Backoff)
					continue
				}

				s.runningMu.Lock()
				s.running--
				s.runningMu.Unlock()
				s.stoppedService <- struct{}{}
				return
			}
		}(name, svc)
	}
	wg.Wait()
}

type backoff struct {
	lastfail time.Time
	failures float64
}

func (b *backoff) wait(failureDecay float64, threshold float64, backoffDur time.Duration) {
	if b.lastfail.IsZero() {
		b.lastfail = time.Now()
		b.failures = 1.0
	} else {
		b.failures++
		intervals := time.Since(b.lastfail).Seconds() / failureDecay
		b.failures = b.failures*math.Pow(.5, intervals) + 1
	}

	if b.failures > threshold {
		time.Sleep(backoffDur)
	}
}
<|endoftext|>


<commit_before>package system

import (
	"fmt"
	"strconv"

	"golang.org/x/net/context"

	"math"

	"kego.io/json"
	"kego.io/kerr"
)

type Int int

func NewInt(i int) *Int {
	out := Int(i)
	return &out
}

func (i *Int) Set(in int) {
	*i = Int(in)
}

func (i *Int) GetString(ctx context.Context) *String {
	return NewString(i.String())
}

func (i *Int) Value() int {
	return int(*i)
}

func (r *IntRule) Enforce(ctx context.Context, data interface{}) (bool, string, error) {

	i, ok := data.(*Int)
	if !ok && data != nil {
		return false, "", kerr.New("AISBHNCJXJ", "Data %T should be *system.Int", data)
	}

	// This provides an upper bound for the restriction
	if r.Maximum != nil {
		if i == nil && !r.Optional {
			return false, "Maximum: value must exist", nil
		}
		if i != nil && i.Value() > r.Maximum.Value() {
			return false, fmt.Sprintf("Maximum: value %v must not be greater than %v", i, r.Maximum.Value()), nil
		}
	}

	// This provides a lower bound for the restriction
	if r.Minimum != nil {
		if i == nil && !r.Optional {
			return false, "Minimum: value must exist", nil
		}
		if i != nil && i.Value() < r.Minimum.Value() {
			return false, fmt.Sprintf("Minimum: value %v must not be less than %v", i, r.Minimum.Value()), nil
		}
	}

	// This restricts the number to be a multiple of the given number
	if r.MultipleOf != nil {
		if i == nil && !r.Optional {
			return false, "MultipleOf: value must exist", nil
		}
		if i != nil && i.Value()%r.MultipleOf.Value() != 0 {
			return false, fmt.Sprintf("MultipleOf: value %v must be a multiple of %v", i, r.MultipleOf.Value()), nil
		}
	}

	return true, "", nil
}

var _ Enforcer = (*IntRule)(nil)

func (out *Int) Unpack(ctx context.Context, in json.Packed) error {
	if in == nil || in.Type() == json.J_NULL {
		return kerr.New("JEJANRWFMH", "Called Int.Unpack with nil value")
	}
	if in.Type() != json.J_NUMBER {
		return kerr.New("UJUBDGVYGF", "Can't unpack %s into *system.Int", in.Type())
	}
	i := math.Floor(in.Number())
	if i != in.Number() {
		return kerr.New("KVEOETSIJY", "%v is not an integer", in.Number())
	}
	*out = Int(int(i))
	return nil
}

var _ json.Unpacker = (*Int)(nil)

func (n *Int) MarshalJSON(ctx context.Context) ([]byte, error) {
	if n == nil {
		return []byte("null"), nil
	}
	return []byte(formatInt(n)), nil
}

var _ json.Marshaler = (*Int)(nil)

func (n *Int) String() string {
	if n == nil {
		return ""
	}
	return formatInt(n)
}

func formatInt(i *Int) string {
	return strconv.FormatInt(int64(*i), 10)
}

func (i Int) NativeNumber() float64 {
	return float64(i)
}

var _ NativeNumber = (*Int)(nil)

func (r *IntRule) GetDefault() interface{} {
	return r.Default
}

var _ DefaultRule = (*IntRule)(nil)
<commit_msg>Set comment<commit_after>package system

import (
	"fmt"
	"strconv"

	"golang.org/x/net/context"

	"math"

	"kego.io/json"
	"kego.io/kerr"
)

type Int int

func NewInt(i int) *Int {
	out := Int(i)
	return &out
}

// Set sets the value
func (i *Int) Set(in int) {
	*i = Int(in)
}

func (i *Int) GetString(ctx context.Context) *String {
	return NewString(i.String())
}

func (i *Int) Value() int {
	return int(*i)
}

func (r *IntRule) Enforce(ctx context.Context, data interface{}) (bool, string, error) {

	i, ok := data.(*Int)
	if !ok && data != nil {
		return false, "", kerr.New("AISBHNCJXJ", "Data %T should be *system.Int", data)
	}

	// This provides an upper bound for the restriction
	if r.Maximum != nil {
		if i == nil && !r.Optional {
			return false, "Maximum: value must exist", nil
		}
		if i != nil && i.Value() > r.Maximum.Value() {
			return false, fmt.Sprintf("Maximum: value %v must not be greater than %v", i, r.Maximum.Value()), nil
		}
	}

	// This provides a lower bound for the restriction
	if r.Minimum != nil {
		if i == nil && !r.Optional {
			return false, "Minimum: value must exist", nil
		}
		if i != nil && i.Value() < r.Minimum.Value() {
			return false, fmt.Sprintf("Minimum: value %v must not be less than %v", i, r.Minimum.Value()), nil
		}
	}

	// This restricts the number to be a multiple of the given number
	if r.MultipleOf != nil {
		if i == nil && !r.Optional {
			return false, "MultipleOf: value must exist", nil
		}
		if i != nil && i.Value()%r.MultipleOf.Value() != 0 {
			return false, fmt.Sprintf("MultipleOf: value %v must be a multiple of %v", i, r.MultipleOf.Value()), nil
		}
	}

	return true, "", nil
}

var _ Enforcer = (*IntRule)(nil)

func (out *Int) Unpack(ctx context.Context, in json.Packed) error {
	if in == nil || in.Type() == json.J_NULL {
		return kerr.New("JEJANRWFMH", "Called Int.Unpack with nil value")
	}
	if in.Type() != json.J_NUMBER {
		return kerr.New("UJUBDGVYGF", "Can't unpack %s into *system.Int", in.Type())
	}
	i := math.Floor(in.Number())
	if i != in.Number() {
		return kerr.New("KVEOETSIJY", "%v is not an integer", in.Number())
	}
	*out = Int(int(i))
	return nil
}

var _ json.Unpacker = (*Int)(nil)

func (n *Int) MarshalJSON(ctx context.Context) ([]byte, error) {
	if n == nil {
		return []byte("null"), nil
	}
	return []byte(formatInt(n)), nil
}

var _ json.Marshaler = (*Int)(nil)

func (n *Int) String() string {
	if n == nil {
		return ""
	}
	return formatInt(n)
}

func formatInt(i *Int) string {
	return strconv.FormatInt(int64(*i), 10)
}

func (i Int) NativeNumber() float64 {
	return float64(i)
}

var _ NativeNumber = (*Int)(nil)

func (r *IntRule) GetDefault() interface{} {
	return r.Default
}

var _ DefaultRule = (*IntRule)(nil)
<|endoftext|>


<commit_before>package target

import (
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"

	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/go.crypto/ssh/agent"
	"github.com/dynport/urknall/cmd"
)

// Create a target for provisioning via SSH.
func NewSshTarget(addr string) (target *sshTarget, e error) {
	target = &sshTarget{port: 22, user: "root"}

	hostAndPort := strings.SplitN(addr, ":", 2)
	if len(hostAndPort) == 2 {
		addr = hostAndPort[0]
		target.port, e = strconv.Atoi(hostAndPort[1])
		if e != nil {
			return nil, fmt.Errorf("port must be given as integer, got %q", hostAndPort[1])
		}
	}

	userAndAddress := strings.Split(addr, "@")
	switch len(userAndAddress) {
	case 1:
		target.address = addr
	case 2:
		target.user = userAndAddress[0]
		target.address = userAndAddress[1]
	default:
		return nil, fmt.Errorf("expected target address of the form '<user>@<host>', but was given: %s", addr)
	}

	if target.address == "" {
		e = fmt.Errorf("empty address given for target")
	}

	return target, e
}

type sshTarget struct {
	Password string

	user    string
	port    int
	address string

	client *ssh.Client
}

func (target *sshTarget) User() string {
	return target.user
}

func (target *sshTarget) String() string {
	return fmt.Sprintf("%s@%s:%d", target.user, target.address, target.port)
}

func (target *sshTarget) Command(cmd string) (cmd.ExecCommand, error) {
	if target.client == nil {
		var e error
		target.client, e = target.buildClient()
		if e != nil {
			return nil, e
		}
	}
	ses, e := target.client.NewSession()
	if e != nil {
		return nil, e
	}
	return &sshCommand{command: cmd, session: ses}, nil
}

func (target *sshTarget) Reset() (e error) {
	if target.client != nil {
		e = target.client.Close()
		target.client = nil
	}
	return e
}

func (target *sshTarget) buildClient() (*ssh.Client, error) {
	var e error
	config := &ssh.ClientConfig{
		User: target.user,
	}
	if target.Password != "" {
		config.Auth = append(config.Auth, ssh.Password(target.Password))
	} else if sshSocket := os.Getenv("SSH_AUTH_SOCK"); sshSocket != "" {
		if agentConn, e := net.Dial("unix", sshSocket); e == nil {
			config.Auth = append(config.Auth, ssh.PublicKeysCallback(agent.NewClient(agentConn).Signers))
		}
	}
	con, e := ssh.Dial("tcp", fmt.Sprintf("%s:%d", target.address, target.port), config)
	if e != nil {
		return nil, e
	}
	return &ssh.Client{Conn: con}, nil
}

type sshCommand struct {
	command string
	session *ssh.Session
}

func (c *sshCommand) Close() error {
	return c.session.Close()
}

func (c *sshCommand) StdinPipe() (io.Writer, error) {
	return c.session.StdinPipe()
}

func (c *sshCommand) StdoutPipe() (io.Reader, error) {
	return c.session.StdoutPipe()
}

func (c *sshCommand) StderrPipe() (io.Reader, error) {
	return c.session.StderrPipe()
}

func (c *sshCommand) SetStdout(w io.Writer) {
	c.session.Stdout = w
}

func (c *sshCommand) SetStderr(w io.Writer) {
	c.session.Stderr = w
}

func (c *sshCommand) SetStdin(r io.Reader) {
	c.session.Stdin = r
}

func (c *sshCommand) Run() error {
	return c.session.Run(c.command)
}

func (c *sshCommand) Wait() error {
	return c.session.Wait()
}

func (c *sshCommand) Start() error {
	return c.session.Start(c.command)
}
<commit_msg>rename code.google.com/p/go.crypto to golang.org/x/crypto/ssh<commit_after>package target

import (
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"

	"github.com/dynport/urknall/cmd"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

// Create a target for provisioning via SSH.
func NewSshTarget(addr string) (target *sshTarget, e error) {
	target = &sshTarget{port: 22, user: "root"}

	hostAndPort := strings.SplitN(addr, ":", 2)
	if len(hostAndPort) == 2 {
		addr = hostAndPort[0]
		target.port, e = strconv.Atoi(hostAndPort[1])
		if e != nil {
			return nil, fmt.Errorf("port must be given as integer, got %q", hostAndPort[1])
		}
	}

	userAndAddress := strings.Split(addr, "@")
	switch len(userAndAddress) {
	case 1:
		target.address = addr
	case 2:
		target.user = userAndAddress[0]
		target.address = userAndAddress[1]
	default:
		return nil, fmt.Errorf("expected target address of the form '<user>@<host>', but was given: %s", addr)
	}

	if target.address == "" {
		e = fmt.Errorf("empty address given for target")
	}

	return target, e
}

type sshTarget struct {
	Password string

	user    string
	port    int
	address string

	client *ssh.Client
}

func (target *sshTarget) User() string {
	return target.user
}

func (target *sshTarget) String() string {
	return fmt.Sprintf("%s@%s:%d", target.user, target.address, target.port)
}

func (target *sshTarget) Command(cmd string) (cmd.ExecCommand, error) {
	if target.client == nil {
		var e error
		target.client, e = target.buildClient()
		if e != nil {
			return nil, e
		}
	}
	ses, e := target.client.NewSession()
	if e != nil {
		return nil, e
	}
	return &sshCommand{command: cmd, session: ses}, nil
}

func (target *sshTarget) Reset() (e error) {
	if target.client != nil {
		e = target.client.Close()
		target.client = nil
	}
	return e
}

func (target *sshTarget) buildClient() (*ssh.Client, error) {
	var e error
	config := &ssh.ClientConfig{
		User: target.user,
	}
	if target.Password != "" {
		config.Auth = append(config.Auth, ssh.Password(target.Password))
	} else if sshSocket := os.Getenv("SSH_AUTH_SOCK"); sshSocket != "" {
		if agentConn, e := net.Dial("unix", sshSocket); e == nil {
			config.Auth = append(config.Auth, ssh.PublicKeysCallback(agent.NewClient(agentConn).Signers))
		}
	}
	con, e := ssh.Dial("tcp", fmt.Sprintf("%s:%d", target.address, target.port), config)
	if e != nil {
		return nil, e
	}
	return &ssh.Client{Conn: con}, nil
}

type sshCommand struct {
	command string
	session *ssh.Session
}

func (c *sshCommand) Close() error {
	return c.session.Close()
}

func (c *sshCommand) StdinPipe() (io.Writer, error) {
	return c.session.StdinPipe()
}

func (c *sshCommand) StdoutPipe() (io.Reader, error) {
	return c.session.StdoutPipe()
}

func (c *sshCommand) StderrPipe() (io.Reader, error) {
	return c.session.StderrPipe()
}

func (c *sshCommand) SetStdout(w io.Writer) {
	c.session.Stdout = w
}

func (c *sshCommand) SetStderr(w io.Writer) {
	c.session.Stderr = w
}

func (c *sshCommand) SetStdin(r io.Reader) {
	c.session.Stdin = r
}

func (c *sshCommand) Run() error {
	return c.session.Run(c.command)
}

func (c *sshCommand) Wait() error {
	return c.session.Wait()
}

func (c *sshCommand) Start() error {
	return c.session.Start(c.command)
}
<|endoftext|>


<commit_before>/*******************************************************************************
*
* Copyright 2018 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, 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 "* 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 reports

import (
	"database/sql"
	"fmt"

	"github.com/sapcc/limes/pkg/db"
	"github.com/sapcc/limes/pkg/limes"
)

//Inconsistencies contains aggregated data about inconsistent quota setups for domains and projects
//in the current cluster.
type Inconsistencies struct {
	ClusterID           string                     `json:"cluster_id"`
	OvercommittedQuotas []OvercommittedDomainQuota `json:"domain_quota_overcommitted,keepempty"`
	// OverspentQuotas     []OverspentProjectQuota    `json:"project_quota_overspent,keepempty"`
	// MismatchQuotas      []MismatchProjectQuota     `json:"project_quota_mismatch,keepempty"`
}

//OvercommittedDomainQuota is a substructure of Inconsistency containing data for the inconsistency type
//where for a domain the sum(projects_quota) > domain_quota for a single resource.
type OvercommittedDomainQuota struct {
	Domain        DomainData `json:"domain,keepempty"`
	Service       string     `json:"service,keepempty"`
	Resource      string     `json:"resource,keepempty"`
	DomainQuota   uint64     `json:"domain_quota,keepempty"`
	ProjectsQuota uint64     `json:"projects_quota,keepempty"`
}

//OverspentProjectQuota is a substructure of Inconsistency containing data for the inconsistency type
//where for some project the usage > quota for a single resource.
// type OverspentProjectQuota struct {
// }

//MismatchProjectQuota is a substructure of Inconsistency containing data for the inconsistency type
//where for some project the quota != backend_quota for a single resource.
// type OverspentProjectQuota struct {
// }

//DomainData is a substructure containing domain data for a single inconsistency
type DomainData struct {
	UUID string `json:"id"`
	Name string `json:"name"`
}

//ProjectData is a substructure containing project data for a single inconsistency
// type ProjectData struct {
// }

var ocdqReportQuery = `
	SELECT d.uuid, d.name, ps.type, pr.name, COALESCE(SUM(dr.quota),0), COALESCE(SUM(pr.quota),0)
	  FROM domains d
	  LEFT OUTER JOIN domain_services ds ON ds.domain_id = d.id
	  LEFT OUTER JOIN domain_resources dr ON dr.service_id = ds.id
	  JOIN projects p ON p.domain_id = d.id
	  LEFT OUTER JOIN project_services ps ON ps.project_id = p.id {{AND ps.type = $service_type}}
	  LEFT OUTER JOIN project_resources pr ON pr.service_id = ps.id {{AND pr.name = $resource_name}}
	WHERE %s GROUP BY d.uuid, d.name, ps.type, pr.name
	HAVING COALESCE(SUM(dr.quota),0) < COALESCE(SUM(pr.quota),0)
	ORDER BY d.uuid ASC
`

//GetInconsistencies returns Inconsistency reports for all inconsistencies and their projects in the current cluster.
func GetInconsistencies(cluster *limes.Cluster, dbi db.Interface, filter Filter) (*Inconsistencies, error) {
	fields := map[string]interface{}{"d.cluster_id": cluster.ID}

	//Initialize inconsistencies as Inconsistencies type and assign ClusterID.
	//The inconsistency data will be assigned in the respective SQL queries.
	inconsistencies := Inconsistencies{
		ClusterID: cluster.ID,
	}

	//ocdqReportQuery: data for OvercommittedDomainQuota inconsistencies.
	queryStr, joinArgs := filter.PrepareQuery(ocdqReportQuery)
	whereStr, whereArgs := db.BuildSimpleWhereClause(fields, len(joinArgs))
	err := db.ForeachRow(db.DB, fmt.Sprintf(queryStr, whereStr), append(joinArgs, whereArgs...), func(rows *sql.Rows) error {
		ocdq := OvercommittedDomainQuota{}
		err := rows.Scan(
			&ocdq.Domain.UUID, &ocdq.Domain.Name, &ocdq.Service, &ocdq.Resource, &ocdq.DomainQuota, &ocdq.ProjectsQuota,
		)
		if err != nil {
			return err
		}

		inconsistencies.OvercommittedQuotas = append(inconsistencies.OvercommittedQuotas, ocdq)

		return err
	})

	if err != nil {
		return nil, err
	}

	return &inconsistencies, nil
}
<commit_msg>add support for project_quota_overspent inconsistencies<commit_after>/*******************************************************************************
*
* Copyright 2018 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, 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 "* 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 reports

import (
	"database/sql"
	"fmt"

	"github.com/sapcc/limes/pkg/db"
	"github.com/sapcc/limes/pkg/limes"
)

//Inconsistencies contains aggregated data about inconsistent quota setups for domains and projects
//in the current cluster.
type Inconsistencies struct {
	ClusterID           string                     `json:"cluster_id"`
	OvercommittedQuotas []OvercommittedDomainQuota `json:"domain_quota_overcommitted,keepempty"`
	OverspentQuotas     []OverspentProjectQuota    `json:"project_quota_overspent,keepempty"`
	// MismatchQuotas      []MismatchProjectQuota     `json:"project_quota_mismatch,keepempty"`
}

//OvercommittedDomainQuota is a substructure of Inconsistency containing data for the inconsistency type
//where for a domain the sum(projects_quota) > domain_quota for a single resource.
type OvercommittedDomainQuota struct {
	Domain        DomainData `json:"domain,keepempty"`
	Service       string     `json:"service,keepempty"`
	Resource      string     `json:"resource,keepempty"`
	DomainQuota   uint64     `json:"domain_quota,keepempty"`
	ProjectsQuota uint64     `json:"projects_quota,keepempty"`
}

//OverspentProjectQuota is a substructure of Inconsistency containing data for the inconsistency type
//where for some project the usage > quota for a single resource.
type OverspentProjectQuota struct {
	Project  ProjectData `json:"project,keepempty"`
	Service  string      `json:"service,keepempty"`
	Resource string      `json:"resource,keepempty"`
	Quota    uint64      `json:"quota,keepempty"`
	Usage    uint64      `json:"usage,keepempty"`
}

//MismatchProjectQuota is a substructure of Inconsistency containing data for the inconsistency type
//where for some project the quota != backend_quota for a single resource.
// type MismatchProjectQuota struct {
// 	Project      ProjectData `json:"project,keepempty"`
// 	Service      string      `json:"service,keepempty"`
// 	Resource     string      `json:"resource,keepempty"`
// 	Quota        uint64      `json:"quota,keepempty"`
// 	BackendQuota uint64      `json:"backend_quota,keepempty"`
// }

//DomainData is a substructure containing domain data for a single inconsistency
type DomainData struct {
	UUID string `json:"id"`
	Name string `json:"name"`
}

//ProjectData is a substructure containing project data for a single inconsistency
type ProjectData struct {
	UUID   string     `json:"id"`
	Name   string     `json:"name"`
	Domain DomainData `json:"domain,keepempty"`
}

var ocdqReportQuery = `
	SELECT d.uuid, d.name, ps.type, pr.name, COALESCE(SUM(dr.quota),0), COALESCE(SUM(pr.quota),0)
	  FROM domains d
	  LEFT OUTER JOIN domain_services ds ON ds.domain_id = d.id
	  LEFT OUTER JOIN domain_resources dr ON dr.service_id = ds.id
	  JOIN projects p ON p.domain_id = d.id
	  LEFT OUTER JOIN project_services ps ON ps.project_id = p.id {{AND ps.type = $service_type}}
	  LEFT OUTER JOIN project_resources pr ON pr.service_id = ps.id {{AND pr.name = $resource_name}}
	WHERE %s GROUP BY d.uuid, d.name, ps.type, pr.name
	HAVING COALESCE(SUM(dr.quota),0) < COALESCE(SUM(pr.quota),0)
	ORDER BY d.uuid ASC
`

var ospqReportQuery = `
	SELECT d.uuid, d.name, p.uuid, p.name, ps.type, pr.name, SUM(pr.quota), SUM(pr.usage)
	  FROM projects p
	  LEFT OUTER JOIN domains d ON d.id=p.domain_id
	  LEFT OUTER JOIN project_services ps ON ps.project_id = p.id {{AND ps.type = $service_type}}
	  LEFT OUTER JOIN project_resources pr ON pr.service_id = ps.id {{AND pr.name = $resource_name}}
	WHERE %s GROUP BY d.uuid, d.name, p.uuid, p.name, ps.type, pr.name
	HAVING SUM(pr.usage) > SUM(pr.quota)
	ORDER BY p.uuid ASC
`

//GetInconsistencies returns Inconsistency reports for all inconsistencies and their projects in the current cluster.
func GetInconsistencies(cluster *limes.Cluster, dbi db.Interface, filter Filter) (*Inconsistencies, error) {
	fields := map[string]interface{}{"d.cluster_id": cluster.ID}

	//Initialize inconsistencies as Inconsistencies type and assign ClusterID.
	//The inconsistency data will be assigned in the respective SQL queries.
	inconsistencies := Inconsistencies{
		ClusterID: cluster.ID,
	}

	//ocdqReportQuery: data for OvercommittedDomainQuota inconsistencies.
	queryStr, joinArgs := filter.PrepareQuery(ocdqReportQuery)
	whereStr, whereArgs := db.BuildSimpleWhereClause(fields, len(joinArgs))
	err := db.ForeachRow(db.DB, fmt.Sprintf(queryStr, whereStr), append(joinArgs, whereArgs...), func(rows *sql.Rows) error {
		ocdq := OvercommittedDomainQuota{}
		err := rows.Scan(
			&ocdq.Domain.UUID, &ocdq.Domain.Name, &ocdq.Service,
			&ocdq.Resource, &ocdq.DomainQuota, &ocdq.ProjectsQuota,
		)
		if err != nil {
			return err
		}

		inconsistencies.OvercommittedQuotas = append(inconsistencies.OvercommittedQuotas, ocdq)

		return err
	})

	//ospqReportQuery: data for OverspentProjectQuota inconsistencies.
	queryStr, joinArgs = filter.PrepareQuery(ospqReportQuery)
	whereStr, whereArgs = db.BuildSimpleWhereClause(fields, len(joinArgs))
	err = db.ForeachRow(db.DB, fmt.Sprintf(queryStr, whereStr), append(joinArgs, whereArgs...), func(rows *sql.Rows) error {
		ospq := OverspentProjectQuota{}
		err := rows.Scan(
			&ospq.Project.Domain.UUID, &ospq.Project.Domain.Name, &ospq.Project.UUID,
			&ospq.Project.Name, &ospq.Service, &ospq.Resource, &ospq.Quota, &ospq.Usage,
		)
		if err != nil {
			return err
		}

		inconsistencies.OverspentQuotas = append(inconsistencies.OverspentQuotas, ospq)

		return err
	})

	if err != nil {
		return nil, err
	}

	return &inconsistencies, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package csi

import (
	"context"
	"errors"
	"fmt"
	"net"
	"time"

	csipb "github.com/container-storage-interface/spec/lib/go/csi/v0"
	"github.com/golang/glog"
	"google.golang.org/grpc"
	api "k8s.io/api/core/v1"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/kubernetes/pkg/features"
)

type csiClient interface {
	NodeGetInfo(ctx context.Context) (
		nodeID string,
		maxVolumePerNode int64,
		accessibleTopology *csipb.Topology,
		err error)
	NodePublishVolume(
		ctx context.Context,
		volumeid string,
		readOnly bool,
		stagingTargetPath string,
		targetPath string,
		accessMode api.PersistentVolumeAccessMode,
		volumeInfo map[string]string,
		volumeAttribs map[string]string,
		nodePublishSecrets map[string]string,
		fsType string,
	) error
	NodeUnpublishVolume(
		ctx context.Context,
		volID string,
		targetPath string,
	) error
	NodeStageVolume(ctx context.Context,
		volID string,
		publishVolumeInfo map[string]string,
		stagingTargetPath string,
		fsType string,
		accessMode api.PersistentVolumeAccessMode,
		nodeStageSecrets map[string]string,
		volumeAttribs map[string]string,
	) error
	NodeUnstageVolume(ctx context.Context, volID, stagingTargetPath string) error
	NodeGetCapabilities(ctx context.Context) ([]*csipb.NodeServiceCapability, error)
}

// csiClient encapsulates all csi-plugin methods
type csiDriverClient struct {
	driverName string
	nodeClient csipb.NodeClient
}

var _ csiClient = &csiDriverClient{}

func newCsiDriverClient(driverName string) *csiDriverClient {
	c := &csiDriverClient{driverName: driverName}
	return c
}

func (c *csiDriverClient) NodeGetInfo(ctx context.Context) (
	nodeID string,
	maxVolumePerNode int64,
	accessibleTopology *csipb.Topology,
	err error) {
	glog.V(4).Info(log("calling NodeGetInfo rpc"))

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return "", 0, nil, err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	res, err := nodeClient.NodeGetInfo(ctx, &csipb.NodeGetInfoRequest{})
	return res.GetNodeId(), res.GetMaxVolumesPerNode(), res.GetAccessibleTopology(), nil
}

func (c *csiDriverClient) NodePublishVolume(
	ctx context.Context,
	volID string,
	readOnly bool,
	stagingTargetPath string,
	targetPath string,
	accessMode api.PersistentVolumeAccessMode,
	volumeInfo map[string]string,
	volumeAttribs map[string]string,
	nodePublishSecrets map[string]string,
	fsType string,
) error {
	glog.V(4).Info(log("calling NodePublishVolume rpc [volid=%s,target_path=%s]", volID, targetPath))
	if volID == "" {
		return errors.New("missing volume id")
	}
	if targetPath == "" {
		return errors.New("missing target path")
	}

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodePublishVolumeRequest{
		VolumeId:           volID,
		TargetPath:         targetPath,
		Readonly:           readOnly,
		PublishInfo:        volumeInfo,
		VolumeAttributes:   volumeAttribs,
		NodePublishSecrets: nodePublishSecrets,
		VolumeCapability: &csipb.VolumeCapability{
			AccessMode: &csipb.VolumeCapability_AccessMode{
				Mode: asCSIAccessMode(accessMode),
			},
		},
	}
	if stagingTargetPath != "" {
		req.StagingTargetPath = stagingTargetPath
	}

	if fsType == fsTypeBlockName {
		req.VolumeCapability.AccessType = &csipb.VolumeCapability_Block{
			Block: &csipb.VolumeCapability_BlockVolume{},
		}
	} else {
		req.VolumeCapability.AccessType = &csipb.VolumeCapability_Mount{
			Mount: &csipb.VolumeCapability_MountVolume{
				FsType: fsType,
			},
		}
	}

	_, err = nodeClient.NodePublishVolume(ctx, req)
	return err
}

func (c *csiDriverClient) NodeUnpublishVolume(ctx context.Context, volID string, targetPath string) error {
	glog.V(4).Info(log("calling NodeUnpublishVolume rpc: [volid=%s, target_path=%s", volID, targetPath))
	if volID == "" {
		return errors.New("missing volume id")
	}
	if targetPath == "" {
		return errors.New("missing target path")
	}

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodeUnpublishVolumeRequest{
		VolumeId:   volID,
		TargetPath: targetPath,
	}

	_, err = nodeClient.NodeUnpublishVolume(ctx, req)
	return err
}

func (c *csiDriverClient) NodeStageVolume(ctx context.Context,
	volID string,
	publishInfo map[string]string,
	stagingTargetPath string,
	fsType string,
	accessMode api.PersistentVolumeAccessMode,
	nodeStageSecrets map[string]string,
	volumeAttribs map[string]string,
) error {
	glog.V(4).Info(log("calling NodeStageVolume rpc [volid=%s,staging_target_path=%s]", volID, stagingTargetPath))
	if volID == "" {
		return errors.New("missing volume id")
	}
	if stagingTargetPath == "" {
		return errors.New("missing staging target path")
	}

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodeStageVolumeRequest{
		VolumeId:          volID,
		PublishInfo:       publishInfo,
		StagingTargetPath: stagingTargetPath,
		VolumeCapability: &csipb.VolumeCapability{
			AccessMode: &csipb.VolumeCapability_AccessMode{
				Mode: asCSIAccessMode(accessMode),
			},
		},
		NodeStageSecrets: nodeStageSecrets,
		VolumeAttributes: volumeAttribs,
	}

	if fsType == fsTypeBlockName {
		req.VolumeCapability.AccessType = &csipb.VolumeCapability_Block{
			Block: &csipb.VolumeCapability_BlockVolume{},
		}
	} else {
		req.VolumeCapability.AccessType = &csipb.VolumeCapability_Mount{
			Mount: &csipb.VolumeCapability_MountVolume{
				FsType: fsType,
			},
		}
	}

	_, err = nodeClient.NodeStageVolume(ctx, req)
	return err
}

func (c *csiDriverClient) NodeUnstageVolume(ctx context.Context, volID, stagingTargetPath string) error {
	glog.V(4).Info(log("calling NodeUnstageVolume rpc [volid=%s,staging_target_path=%s]", volID, stagingTargetPath))
	if volID == "" {
		return errors.New("missing volume id")
	}
	if stagingTargetPath == "" {
		return errors.New("missing staging target path")
	}

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodeUnstageVolumeRequest{
		VolumeId:          volID,
		StagingTargetPath: stagingTargetPath,
	}
	_, err = nodeClient.NodeUnstageVolume(ctx, req)
	return err
}

func (c *csiDriverClient) NodeGetCapabilities(ctx context.Context) ([]*csipb.NodeServiceCapability, error) {
	glog.V(4).Info(log("calling NodeGetCapabilities rpc"))

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodeGetCapabilitiesRequest{}
	resp, err := nodeClient.NodeGetCapabilities(ctx, req)
	if err != nil {
		return nil, err
	}
	return resp.GetCapabilities(), nil
}

func asCSIAccessMode(am api.PersistentVolumeAccessMode) csipb.VolumeCapability_AccessMode_Mode {
	switch am {
	case api.ReadWriteOnce:
		return csipb.VolumeCapability_AccessMode_SINGLE_NODE_WRITER
	case api.ReadOnlyMany:
		return csipb.VolumeCapability_AccessMode_MULTI_NODE_READER_ONLY
	case api.ReadWriteMany:
		return csipb.VolumeCapability_AccessMode_MULTI_NODE_MULTI_WRITER
	}
	return csipb.VolumeCapability_AccessMode_UNKNOWN
}

func newGrpcConn(driverName string) (*grpc.ClientConn, error) {
	if driverName == "" {
		return nil, fmt.Errorf("driver name is empty")
	}
	addr := fmt.Sprintf(csiAddrTemplate, driverName)
	// TODO once KubeletPluginsWatcher graduates to beta, remove FeatureGate check
	if utilfeature.DefaultFeatureGate.Enabled(features.KubeletPluginsWatcher) {
		driver, ok := csiDrivers.driversMap[driverName]
		if !ok {
			return nil, fmt.Errorf("driver name %s not found in the list of registered CSI drivers", driverName)
		}
		addr = driver.driverEndpoint
	}
	network := "unix"
	glog.V(4).Infof(log("creating new gRPC connection for [%s://%s]", network, addr))

	return grpc.Dial(
		addr,
		grpc.WithInsecure(),
		grpc.WithDialer(func(target string, timeout time.Duration) (net.Conn, error) {
			return net.Dial(network, target)
		}),
	)
}
<commit_msg>Return error from NodeGetInfo<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package csi

import (
	"context"
	"errors"
	"fmt"
	"net"
	"time"

	csipb "github.com/container-storage-interface/spec/lib/go/csi/v0"
	"github.com/golang/glog"
	"google.golang.org/grpc"
	api "k8s.io/api/core/v1"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/kubernetes/pkg/features"
)

type csiClient interface {
	NodeGetInfo(ctx context.Context) (
		nodeID string,
		maxVolumePerNode int64,
		accessibleTopology *csipb.Topology,
		err error)
	NodePublishVolume(
		ctx context.Context,
		volumeid string,
		readOnly bool,
		stagingTargetPath string,
		targetPath string,
		accessMode api.PersistentVolumeAccessMode,
		volumeInfo map[string]string,
		volumeAttribs map[string]string,
		nodePublishSecrets map[string]string,
		fsType string,
	) error
	NodeUnpublishVolume(
		ctx context.Context,
		volID string,
		targetPath string,
	) error
	NodeStageVolume(ctx context.Context,
		volID string,
		publishVolumeInfo map[string]string,
		stagingTargetPath string,
		fsType string,
		accessMode api.PersistentVolumeAccessMode,
		nodeStageSecrets map[string]string,
		volumeAttribs map[string]string,
	) error
	NodeUnstageVolume(ctx context.Context, volID, stagingTargetPath string) error
	NodeGetCapabilities(ctx context.Context) ([]*csipb.NodeServiceCapability, error)
}

// csiClient encapsulates all csi-plugin methods
type csiDriverClient struct {
	driverName string
	nodeClient csipb.NodeClient
}

var _ csiClient = &csiDriverClient{}

func newCsiDriverClient(driverName string) *csiDriverClient {
	c := &csiDriverClient{driverName: driverName}
	return c
}

func (c *csiDriverClient) NodeGetInfo(ctx context.Context) (
	nodeID string,
	maxVolumePerNode int64,
	accessibleTopology *csipb.Topology,
	err error) {
	glog.V(4).Info(log("calling NodeGetInfo rpc"))

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return "", 0, nil, err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	res, err := nodeClient.NodeGetInfo(ctx, &csipb.NodeGetInfoRequest{})
	if err != nil {
		return "", 0, nil, err
	}

	return res.GetNodeId(), res.GetMaxVolumesPerNode(), res.GetAccessibleTopology(), nil
}

func (c *csiDriverClient) NodePublishVolume(
	ctx context.Context,
	volID string,
	readOnly bool,
	stagingTargetPath string,
	targetPath string,
	accessMode api.PersistentVolumeAccessMode,
	volumeInfo map[string]string,
	volumeAttribs map[string]string,
	nodePublishSecrets map[string]string,
	fsType string,
) error {
	glog.V(4).Info(log("calling NodePublishVolume rpc [volid=%s,target_path=%s]", volID, targetPath))
	if volID == "" {
		return errors.New("missing volume id")
	}
	if targetPath == "" {
		return errors.New("missing target path")
	}

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodePublishVolumeRequest{
		VolumeId:           volID,
		TargetPath:         targetPath,
		Readonly:           readOnly,
		PublishInfo:        volumeInfo,
		VolumeAttributes:   volumeAttribs,
		NodePublishSecrets: nodePublishSecrets,
		VolumeCapability: &csipb.VolumeCapability{
			AccessMode: &csipb.VolumeCapability_AccessMode{
				Mode: asCSIAccessMode(accessMode),
			},
		},
	}
	if stagingTargetPath != "" {
		req.StagingTargetPath = stagingTargetPath
	}

	if fsType == fsTypeBlockName {
		req.VolumeCapability.AccessType = &csipb.VolumeCapability_Block{
			Block: &csipb.VolumeCapability_BlockVolume{},
		}
	} else {
		req.VolumeCapability.AccessType = &csipb.VolumeCapability_Mount{
			Mount: &csipb.VolumeCapability_MountVolume{
				FsType: fsType,
			},
		}
	}

	_, err = nodeClient.NodePublishVolume(ctx, req)
	return err
}

func (c *csiDriverClient) NodeUnpublishVolume(ctx context.Context, volID string, targetPath string) error {
	glog.V(4).Info(log("calling NodeUnpublishVolume rpc: [volid=%s, target_path=%s", volID, targetPath))
	if volID == "" {
		return errors.New("missing volume id")
	}
	if targetPath == "" {
		return errors.New("missing target path")
	}

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodeUnpublishVolumeRequest{
		VolumeId:   volID,
		TargetPath: targetPath,
	}

	_, err = nodeClient.NodeUnpublishVolume(ctx, req)
	return err
}

func (c *csiDriverClient) NodeStageVolume(ctx context.Context,
	volID string,
	publishInfo map[string]string,
	stagingTargetPath string,
	fsType string,
	accessMode api.PersistentVolumeAccessMode,
	nodeStageSecrets map[string]string,
	volumeAttribs map[string]string,
) error {
	glog.V(4).Info(log("calling NodeStageVolume rpc [volid=%s,staging_target_path=%s]", volID, stagingTargetPath))
	if volID == "" {
		return errors.New("missing volume id")
	}
	if stagingTargetPath == "" {
		return errors.New("missing staging target path")
	}

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodeStageVolumeRequest{
		VolumeId:          volID,
		PublishInfo:       publishInfo,
		StagingTargetPath: stagingTargetPath,
		VolumeCapability: &csipb.VolumeCapability{
			AccessMode: &csipb.VolumeCapability_AccessMode{
				Mode: asCSIAccessMode(accessMode),
			},
		},
		NodeStageSecrets: nodeStageSecrets,
		VolumeAttributes: volumeAttribs,
	}

	if fsType == fsTypeBlockName {
		req.VolumeCapability.AccessType = &csipb.VolumeCapability_Block{
			Block: &csipb.VolumeCapability_BlockVolume{},
		}
	} else {
		req.VolumeCapability.AccessType = &csipb.VolumeCapability_Mount{
			Mount: &csipb.VolumeCapability_MountVolume{
				FsType: fsType,
			},
		}
	}

	_, err = nodeClient.NodeStageVolume(ctx, req)
	return err
}

func (c *csiDriverClient) NodeUnstageVolume(ctx context.Context, volID, stagingTargetPath string) error {
	glog.V(4).Info(log("calling NodeUnstageVolume rpc [volid=%s,staging_target_path=%s]", volID, stagingTargetPath))
	if volID == "" {
		return errors.New("missing volume id")
	}
	if stagingTargetPath == "" {
		return errors.New("missing staging target path")
	}

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodeUnstageVolumeRequest{
		VolumeId:          volID,
		StagingTargetPath: stagingTargetPath,
	}
	_, err = nodeClient.NodeUnstageVolume(ctx, req)
	return err
}

func (c *csiDriverClient) NodeGetCapabilities(ctx context.Context) ([]*csipb.NodeServiceCapability, error) {
	glog.V(4).Info(log("calling NodeGetCapabilities rpc"))

	conn, err := newGrpcConn(c.driverName)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	nodeClient := csipb.NewNodeClient(conn)

	req := &csipb.NodeGetCapabilitiesRequest{}
	resp, err := nodeClient.NodeGetCapabilities(ctx, req)
	if err != nil {
		return nil, err
	}
	return resp.GetCapabilities(), nil
}

func asCSIAccessMode(am api.PersistentVolumeAccessMode) csipb.VolumeCapability_AccessMode_Mode {
	switch am {
	case api.ReadWriteOnce:
		return csipb.VolumeCapability_AccessMode_SINGLE_NODE_WRITER
	case api.ReadOnlyMany:
		return csipb.VolumeCapability_AccessMode_MULTI_NODE_READER_ONLY
	case api.ReadWriteMany:
		return csipb.VolumeCapability_AccessMode_MULTI_NODE_MULTI_WRITER
	}
	return csipb.VolumeCapability_AccessMode_UNKNOWN
}

func newGrpcConn(driverName string) (*grpc.ClientConn, error) {
	if driverName == "" {
		return nil, fmt.Errorf("driver name is empty")
	}
	addr := fmt.Sprintf(csiAddrTemplate, driverName)
	// TODO once KubeletPluginsWatcher graduates to beta, remove FeatureGate check
	if utilfeature.DefaultFeatureGate.Enabled(features.KubeletPluginsWatcher) {
		driver, ok := csiDrivers.driversMap[driverName]
		if !ok {
			return nil, fmt.Errorf("driver name %s not found in the list of registered CSI drivers", driverName)
		}
		addr = driver.driverEndpoint
	}
	network := "unix"
	glog.V(4).Infof(log("creating new gRPC connection for [%s://%s]", network, addr))

	return grpc.Dial(
		addr,
		grpc.WithInsecure(),
		grpc.WithDialer(func(target string, timeout time.Duration) (net.Conn, error) {
			return net.Dial(network, target)
		}),
	)
}
<|endoftext|>


<commit_before>package werf

import (
	"context"
	"fmt"
	"path/filepath"
	"time"

	"github.com/werf/lockgate"
	"github.com/werf/werf/pkg/util/timestamps"
)

func getWerfLastRunAtPath() string {
	return filepath.Join(GetServiceDir(), "var", "last_werf_run_at")
}

func getWerfFirstRunAtPath() string {
	return filepath.Join(GetServiceDir(), "var", "first_werf_run_at")
}

func SetWerfLastRunAt(ctx context.Context) error {
	path := getWerfLastRunAtPath()
	if _, lock, err := AcquireHostLock(ctx, path, lockgate.AcquireOptions{OnWaitFunc: func(lockName string, doWait func() error) error { return doWait() }}); err != nil {
		return fmt.Errorf("error locking path %q: %s", path, err)
	} else {
		defer ReleaseHostLock(lock)
	}

	return timestamps.WriteTimestampFile(path, time.Now())
}

func GetWerfLastRunAt(ctx context.Context) (time.Time, error) {
	path := getWerfLastRunAtPath()
	if _, lock, err := AcquireHostLock(ctx, path, lockgate.AcquireOptions{OnWaitFunc: func(lockName string, doWait func() error) error { return doWait() }}); err != nil {
		return time.Time{}, fmt.Errorf("error locking path %q: %s", path, err)
	} else {
		defer ReleaseHostLock(lock)
	}

	return timestamps.ReadTimestampFile(path)
}

func SetWerfFirstRunAt(ctx context.Context) error {
	path := getWerfFirstRunAtPath()
	if _, lock, err := AcquireHostLock(ctx, path, lockgate.AcquireOptions{OnWaitFunc: func(lockName string, doWait func() error) error { return doWait() }}); err != nil {
		return fmt.Errorf("error locking path %q: %s", path, err)
	} else {
		defer ReleaseHostLock(lock)
	}

	if exists, err := timestamps.CheckTimestampFileExists(path); err != nil {
		return fmt.Errorf("error checking existance of %q: %s", path, err)
	} else if !exists {
		return timestamps.WriteTimestampFile(path, time.Now())
	}
	return nil
}

func GetWerfFirstRunAt(ctx context.Context) (time.Time, error) {
	path := getWerfFirstRunAtPath()
	if _, lock, err := AcquireHostLock(ctx, path, lockgate.AcquireOptions{OnWaitFunc: func(lockName string, doWait func() error) error { return doWait() }}); err != nil {
		return time.Time{}, fmt.Errorf("error locking path %q: %s", path, err)
	} else {
		defer ReleaseHostLock(lock)
	}

	return timestamps.ReadTimestampFile(path)
}
<commit_msg>[GC] Change service werf info dir to ~/.werf/service/info/{first_werf_run_at|last_werf_run_at}<commit_after>package werf

import (
	"context"
	"fmt"
	"path/filepath"
	"time"

	"github.com/werf/lockgate"
	"github.com/werf/werf/pkg/util/timestamps"
)

func getWerfLastRunAtPath() string {
	return filepath.Join(GetServiceDir(), "info", "v1.2", "last_werf_run_at")
}

func getWerfFirstRunAtPath() string {
	return filepath.Join(GetServiceDir(), "info", "v1.2", "first_werf_run_at")
}

func SetWerfLastRunAt(ctx context.Context) error {
	path := getWerfLastRunAtPath()
	if _, lock, err := AcquireHostLock(ctx, path, lockgate.AcquireOptions{OnWaitFunc: func(lockName string, doWait func() error) error { return doWait() }}); err != nil {
		return fmt.Errorf("error locking path %q: %s", path, err)
	} else {
		defer ReleaseHostLock(lock)
	}

	return timestamps.WriteTimestampFile(path, time.Now())
}

func GetWerfLastRunAt(ctx context.Context) (time.Time, error) {
	path := getWerfLastRunAtPath()
	if _, lock, err := AcquireHostLock(ctx, path, lockgate.AcquireOptions{OnWaitFunc: func(lockName string, doWait func() error) error { return doWait() }}); err != nil {
		return time.Time{}, fmt.Errorf("error locking path %q: %s", path, err)
	} else {
		defer ReleaseHostLock(lock)
	}

	return timestamps.ReadTimestampFile(path)
}

func SetWerfFirstRunAt(ctx context.Context) error {
	path := getWerfFirstRunAtPath()
	if _, lock, err := AcquireHostLock(ctx, path, lockgate.AcquireOptions{OnWaitFunc: func(lockName string, doWait func() error) error { return doWait() }}); err != nil {
		return fmt.Errorf("error locking path %q: %s", path, err)
	} else {
		defer ReleaseHostLock(lock)
	}

	if exists, err := timestamps.CheckTimestampFileExists(path); err != nil {
		return fmt.Errorf("error checking existance of %q: %s", path, err)
	} else if !exists {
		return timestamps.WriteTimestampFile(path, time.Now())
	}
	return nil
}

func GetWerfFirstRunAt(ctx context.Context) (time.Time, error) {
	path := getWerfFirstRunAtPath()
	if _, lock, err := AcquireHostLock(ctx, path, lockgate.AcquireOptions{OnWaitFunc: func(lockName string, doWait func() error) error { return doWait() }}); err != nil {
		return time.Time{}, fmt.Errorf("error locking path %q: %s", path, err)
	} else {
		defer ReleaseHostLock(lock)
	}

	return timestamps.ReadTimestampFile(path)
}
<|endoftext|>