<commit_before>package api

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

	"github.com/FederationOfFathers/dashboard/bot"
	"github.com/FederationOfFathers/dashboard/bridge"
	"github.com/FederationOfFathers/dashboard/store"
	"github.com/gorilla/mux"
	"github.com/nlopes/slack"
	"github.com/uber-go/zap"
	stow "gopkg.in/djherbis/stow.v2"
)

func visDB() *stow.Store {
	return store.DB.Groups().NewNestedStore([]byte("visibility-v1"))
}

func isSlackUserInGroup(slackID, groupID string) bool {
	for _, group := range bridge.Data.Slack.GetGroups() {
		if group.ID != groupID {
			continue
		}
		for _, member := range group.Members {
			if member == slackID {
				return true
			}
		}
		return false
	}
	return false
}

func init() {
	Router.Path("/api/v0/groups").Methods("GET").Handler(jwtHandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			id := getSlackUserID(r)
			admin, _ := bridge.Data.Slack.IsUserIDAdmin(id)
			var groups = map[string]map[string]interface{}{}
			for _, group := range bridge.Data.Slack.GetGroups() {
				var visible string
				visDB().Get(group.ID, &visible)
				if visible != "true" {
					visible = "false"
				}
				if visible == "false" {
					if !admin {
						continue
					}
				}
				members := []string{}
				for _, memberID := range group.Members {
					user, err := bridge.Data.Slack.User(memberID)
					if err != nil {
						continue
					}
					members = append(members, user.Name)
				}
				groups[group.ID] = map[string]interface{}{
					"id":      group.ID,
					"name":    group.Name,
					"topic":   group.Topic.Value,
					"purpose": group.Purpose.Value,
					"members": members,
					"visible": visible,
				}
			}
			enc := json.NewEncoder(w)
			enc.Encode(groups)
		},
	))

	Router.Path("/api/v0/groups/{groupID}/join").Methods("GET").Handler(jwtHandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			id := getSlackUserID(r)
			want := mux.Vars(r)["groupID"]
			var visible string
			visDB().Get(want, &visible)
			if visible == "true" || isSlackUserInGroup(id, want) {
				for _, group := range bridge.Data.Slack.GetGroups() {
					if group.ID != want {
						continue
					}
					if err := bot.GroupInvite(group.ID, id); err != nil {
						w.WriteHeader(http.StatusInternalServerError)
					}
					return
				}
			}
			w.WriteHeader(http.StatusNotFound)
		},
	))

	Router.Path("/api/v0/groups/{groupID}/leave").Methods("GET").Handler(jwtHandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			id := getSlackUserID(r)
			want := mux.Vars(r)["groupID"]
			var visible string
			visDB().Get(want, &visible)
			if visible == "true" || isSlackUserInGroup(id, want) {
				for _, group := range bridge.Data.Slack.GetGroups() {
					if group.ID != want {
						continue
					}
					if err := bot.GroupKick(group.ID, id); err != nil {
						w.WriteHeader(http.StatusInternalServerError)
					}
					return
				}
			}
			w.WriteHeader(http.StatusNotFound)
		},
	))

	Router.Path("/api/v0/groups/{groupID}/visibility").Methods("PUT", "POST", "OPTIONS").Handler(jwtHandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			defer r.Body.Close()
			if r.Method == "OPTIONS" {
				return
			}

			var group *slack.Group

			w.Header().Set("Content-Type", "application/json")

			for _, g := range bridge.Data.Slack.GetGroups() {
				if mux.Vars(r)["groupID"] == g.ID {
					group = &g
					break
				}
			}

			if group == nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			id := getSlackUserID(r)
			member, err := DB.MemberBySlackID(id)

			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			if err := requireAdmin(w, r); err != nil {
				// Not an admin... so a request
				bot.SendMessage(
					"damnbot-admin",
					fmt.Sprintf(
						"Request to change group visibility for *%s* received from *%s*",
						group.Name,
						member.Name,
					),
				)
				w.WriteHeader(http.StatusOK)
				return
			}

			var doc = struct {
				Visibility string `json:"visible"`
			}{}

			if err := json.NewDecoder(r.Body).Decode(&doc); err != nil {
				logger.Error("Failed decoding json for group visibility change", zap.Error(err))
				w.WriteHeader(http.StatusBadRequest)
				return
			}

			if err = visDB().Put(mux.Vars(r)["groupID"], doc.Visibility); err != nil {
				logger.Error("error putting a value to visDB", zap.Error(err))
				w.WriteHeader(http.StatusInternalServerError)
			}
		},
	))
}
<commit_msg>200 ok for options always for visibility<commit_after>package api

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

	"github.com/FederationOfFathers/dashboard/bot"
	"github.com/FederationOfFathers/dashboard/bridge"
	"github.com/FederationOfFathers/dashboard/store"
	"github.com/gorilla/mux"
	"github.com/nlopes/slack"
	"github.com/uber-go/zap"
	stow "gopkg.in/djherbis/stow.v2"
)

func visDB() *stow.Store {
	return store.DB.Groups().NewNestedStore([]byte("visibility-v1"))
}

func isSlackUserInGroup(slackID, groupID string) bool {
	for _, group := range bridge.Data.Slack.GetGroups() {
		if group.ID != groupID {
			continue
		}
		for _, member := range group.Members {
			if member == slackID {
				return true
			}
		}
		return false
	}
	return false
}

func init() {
	Router.Path("/api/v0/groups").Methods("GET").Handler(jwtHandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			id := getSlackUserID(r)
			admin, _ := bridge.Data.Slack.IsUserIDAdmin(id)
			var groups = map[string]map[string]interface{}{}
			for _, group := range bridge.Data.Slack.GetGroups() {
				var visible string
				visDB().Get(group.ID, &visible)
				if visible != "true" {
					visible = "false"
				}
				if visible == "false" {
					if !admin {
						continue
					}
				}
				members := []string{}
				for _, memberID := range group.Members {
					user, err := bridge.Data.Slack.User(memberID)
					if err != nil {
						continue
					}
					members = append(members, user.Name)
				}
				groups[group.ID] = map[string]interface{}{
					"id":      group.ID,
					"name":    group.Name,
					"topic":   group.Topic.Value,
					"purpose": group.Purpose.Value,
					"members": members,
					"visible": visible,
				}
			}
			enc := json.NewEncoder(w)
			enc.Encode(groups)
		},
	))

	Router.Path("/api/v0/groups/{groupID}/join").Methods("GET").Handler(jwtHandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			id := getSlackUserID(r)
			want := mux.Vars(r)["groupID"]
			var visible string
			visDB().Get(want, &visible)
			if visible == "true" || isSlackUserInGroup(id, want) {
				for _, group := range bridge.Data.Slack.GetGroups() {
					if group.ID != want {
						continue
					}
					if err := bot.GroupInvite(group.ID, id); err != nil {
						w.WriteHeader(http.StatusInternalServerError)
					}
					return
				}
			}
			w.WriteHeader(http.StatusNotFound)
		},
	))

	Router.Path("/api/v0/groups/{groupID}/leave").Methods("GET").Handler(jwtHandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			id := getSlackUserID(r)
			want := mux.Vars(r)["groupID"]
			var visible string
			visDB().Get(want, &visible)
			if visible == "true" || isSlackUserInGroup(id, want) {
				for _, group := range bridge.Data.Slack.GetGroups() {
					if group.ID != want {
						continue
					}
					if err := bot.GroupKick(group.ID, id); err != nil {
						w.WriteHeader(http.StatusInternalServerError)
					}
					return
				}
			}
			w.WriteHeader(http.StatusNotFound)
		},
	))

	Router.Path("/api/v0/groups/{groupID}/visibility").Methods("PUT", "POST", "OPTIONS").Handler(jwtHandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			defer r.Body.Close()
			if r.Method == "OPTIONS" {
				w.WriteHeader(http.StatusOK)
				return
			}

			var group *slack.Group

			w.Header().Set("Content-Type", "application/json")

			for _, g := range bridge.Data.Slack.GetGroups() {
				if mux.Vars(r)["groupID"] == g.ID {
					group = &g
					break
				}
			}

			if group == nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			id := getSlackUserID(r)
			member, err := DB.MemberBySlackID(id)

			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			if err := requireAdmin(w, r); err != nil {
				// Not an admin... so a request
				bot.SendMessage(
					"damnbot-admin",
					fmt.Sprintf(
						"Request to change group visibility for *%s* received from *%s*",
						group.Name,
						member.Name,
					),
				)
				w.WriteHeader(http.StatusOK)
				return
			}

			var doc = struct {
				Visibility string `json:"visible"`
			}{}

			if err := json.NewDecoder(r.Body).Decode(&doc); err != nil {
				logger.Error("Failed decoding json for group visibility change", zap.Error(err))
				w.WriteHeader(http.StatusBadRequest)
				return
			}

			if err = visDB().Put(mux.Vars(r)["groupID"], doc.Visibility); err != nil {
				logger.Error("error putting a value to visDB", zap.Error(err))
				w.WriteHeader(http.StatusInternalServerError)
			}
		},
	))
}
<|endoftext|>


<commit_before>package solver

import "fmt"

// A Clause is a list of Lit, associated with possible data (for learned clauses).
type Clause struct {
	lits     []Lit
	lbdValue uint32 // contains value for lbd (30 lowest bytes) but also flags to indicate whether clause is learned or not, and locked or not
	activity float32
}

const (
	learnedMask uint32 = 1 << 31
	lockedMask  uint32 = 1 << 30
	bothMasks   uint32 = learnedMask | lockedMask
)

// NewClause returns a clause whose lits are given as an argument.
func NewClause(lits []Lit) *Clause {
	return &Clause{lits: lits}
}

// NewLearnedClause returns a new clause marked as learned.
func NewLearnedClause(lits []Lit) *Clause {
	return &Clause{lits: lits, lbdValue: learnedMask}
}

// Learned returns true iff c was a learned clause.
func (c *Clause) Learned() bool {
	return c.lbdValue&learnedMask == learnedMask
}

func (c *Clause) lock() {
	// c.locked = true
	c.lbdValue = c.lbdValue | lockedMask
}

func (c *Clause) unlock() {
	// c.locked = false
	c.lbdValue = c.lbdValue & ^lockedMask
}

func (c *Clause) lbd() int {
	return int(c.lbdValue & ^bothMasks)
}

func (c *Clause) setLbd(lbd int) {
	c.lbdValue = (c.lbdValue & bothMasks) | uint32(lbd)
}

func (c *Clause) incLbd() {
	c.lbdValue++
}

func (c *Clause) isLocked() bool {
	// return c.learned && c.locked
	return c.lbdValue&bothMasks == bothMasks
}

// Len returns the nb of lits in the clause.
func (c *Clause) Len() int {
	return len(c.lits)
}

// First returns the first lit from the clause.
func (c *Clause) First() Lit {
	return c.lits[0]
}

// Second returns the second lit from the clause.
func (c *Clause) Second() Lit {
	return c.lits[1]
}

// Get returns the ith literal from the clause.
func (c *Clause) Get(i int) Lit {
	return c.lits[i]
}

// Set sets the ith literal of the clause.
func (c *Clause) Set(i int, l Lit) {
	c.lits[i] = l
}

// swapFirstWith swaps the first and ith lits from the clause.
func (c *Clause) swapFirstWith(i int) {
	c.lits[0], c.lits[i] = c.lits[i], c.lits[0]
}

// swapSecondWith swaps the second and ith lits from the clause.
func (c *Clause) swapSecondWith(i int) {
	c.lits[1], c.lits[i] = c.lits[i], c.lits[1]
}

// Shrink reduces the length of the clauses, by removing all lits
// starting from position newLen.
func (c *Clause) Shrink(newLen int) {
	c.lits = c.lits[:newLen]
}

// CNF returns a DIMACS CNF representation of the clause.
func (c *Clause) CNF() string {
	res := ""
	for _, lit := range c.lits {
		res += fmt.Sprintf("%d ", lit.Int())
	}
	return fmt.Sprintf("%s0", res)
}

// OutputClause displays a clause on stdout.
func OutputClause(c *Clause) {
	fmt.Printf("[")
	for i, l := range c.lits {
		if i < c.Len()-1 {
			fmt.Printf("%d, ", l.Int())
		} else {
			fmt.Printf("%d", l.Int())
		}
	}
	fmt.Printf("]\n")
}
<commit_msg>removed dead code in comments<commit_after>package solver

import "fmt"

// A Clause is a list of Lit, associated with possible data (for learned clauses).
type Clause struct {
	lits     []Lit
	lbdValue uint32 // contains value for lbd (30 lowest bytes) but also flags to indicate whether clause is learned or not, and locked or not
	activity float32
}

const (
	learnedMask uint32 = 1 << 31
	lockedMask  uint32 = 1 << 30
	bothMasks   uint32 = learnedMask | lockedMask
)

// NewClause returns a clause whose lits are given as an argument.
func NewClause(lits []Lit) *Clause {
	return &Clause{lits: lits}
}

// NewLearnedClause returns a new clause marked as learned.
func NewLearnedClause(lits []Lit) *Clause {
	return &Clause{lits: lits, lbdValue: learnedMask}
}

// Learned returns true iff c was a learned clause.
func (c *Clause) Learned() bool {
	return c.lbdValue&learnedMask == learnedMask
}

func (c *Clause) lock() {
	c.lbdValue = c.lbdValue | lockedMask
}

func (c *Clause) unlock() {
	c.lbdValue = c.lbdValue & ^lockedMask
}

func (c *Clause) lbd() int {
	return int(c.lbdValue & ^bothMasks)
}

func (c *Clause) setLbd(lbd int) {
	c.lbdValue = (c.lbdValue & bothMasks) | uint32(lbd)
}

func (c *Clause) incLbd() {
	c.lbdValue++
}

func (c *Clause) isLocked() bool {
	return c.lbdValue&bothMasks == bothMasks
}

// Len returns the nb of lits in the clause.
func (c *Clause) Len() int {
	return len(c.lits)
}

// First returns the first lit from the clause.
func (c *Clause) First() Lit {
	return c.lits[0]
}

// Second returns the second lit from the clause.
func (c *Clause) Second() Lit {
	return c.lits[1]
}

// Get returns the ith literal from the clause.
func (c *Clause) Get(i int) Lit {
	return c.lits[i]
}

// Set sets the ith literal of the clause.
func (c *Clause) Set(i int, l Lit) {
	c.lits[i] = l
}

// swapFirstWith swaps the first and ith lits from the clause.
func (c *Clause) swapFirstWith(i int) {
	c.lits[0], c.lits[i] = c.lits[i], c.lits[0]
}

// swapSecondWith swaps the second and ith lits from the clause.
func (c *Clause) swapSecondWith(i int) {
	c.lits[1], c.lits[i] = c.lits[i], c.lits[1]
}

// Shrink reduces the length of the clauses, by removing all lits
// starting from position newLen.
func (c *Clause) Shrink(newLen int) {
	c.lits = c.lits[:newLen]
}

// CNF returns a DIMACS CNF representation of the clause.
func (c *Clause) CNF() string {
	res := ""
	for _, lit := range c.lits {
		res += fmt.Sprintf("%d ", lit.Int())
	}
	return fmt.Sprintf("%s0", res)
}

// OutputClause displays a clause on stdout.
func OutputClause(c *Clause) {
	fmt.Printf("[")
	for i, l := range c.lits {
		if i < c.Len()-1 {
			fmt.Printf("%d, ", l.Int())
		} else {
			fmt.Printf("%d", l.Int())
		}
	}
	fmt.Printf("]\n")
}
<|endoftext|>


<commit_before>package commands

import (
	"bytes"
	"io"
	"strings"
	"sync"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"

	cmds "github.com/jbenet/go-ipfs/commands"
	"github.com/jbenet/go-ipfs/core"
	dag "github.com/jbenet/go-ipfs/merkledag"
	path "github.com/jbenet/go-ipfs/path"
	u "github.com/jbenet/go-ipfs/util"
)

// KeyList is a general type for outputting lists of keys
type KeyList struct {
	Keys []u.Key
}

// KeyListTextMarshaler outputs a KeyList as plaintext, one key per line
func KeyListTextMarshaler(res cmds.Response) (io.Reader, error) {
	output := res.Output().(*KeyList)
	var buf bytes.Buffer
	for _, key := range output.Keys {
		buf.WriteString(key.B58String() + "\n")
	}
	return &buf, nil
}

var RefsCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists links (references) from an object",
		ShortDescription: `
Retrieves the object named by <ipfs-path> and displays the link
hashes it contains, with the following format:

  <link base58 hash>

Note: list all refs recursively with -r.
`,
	},
	Subcommands: map[string]*cmds.Command{
		"local": RefsLocalCmd,
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("ipfs-path", true, true, "Path to the object(s) to list refs from").EnableStdin(),
	},
	Options: []cmds.Option{
		cmds.StringOption("format", "Emit edges with given format. tokens: <src> <dst> <linkname>"),
		cmds.BoolOption("edges", "e", "Emit edge format: `<from> -> <to>`"),
		cmds.BoolOption("unique", "u", "Omit duplicate refs from output"),
		cmds.BoolOption("recursive", "r", "Recursively list links of child nodes"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context().Context
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		unique, _, err := req.Option("unique").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		recursive, _, err := req.Option("recursive").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		edges, _, err := req.Option("edges").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		format, _, err := req.Option("format").String()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		objs, err := objectsForPaths(n, req.Arguments())
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		piper, pipew := io.Pipe()
		eptr := &ErrPassThroughReader{R: piper}

		go func() {
			defer pipew.Close()

			rw := RefWriter{
				W:         pipew,
				DAG:       n.DAG,
				Ctx:       ctx,
				Unique:    unique,
				PrintEdge: edges,
				PrintFmt:  format,
				Recursive: recursive,
			}

			for _, o := range objs {
				if _, err := rw.WriteRefs(o); err != nil {
					eptr.SetError(err)
					return
				}
			}
		}()

		res.SetOutput(eptr)
	},
}

var RefsLocalCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists all local references",
		ShortDescription: `
Displays the hashes of all local objects.
`,
	},

	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context().Context
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		// todo: make async
		allKeys, err := n.Blockstore.AllKeysChan(ctx)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		piper, pipew := io.Pipe()
		eptr := &ErrPassThroughReader{R: piper}

		go func() {
			defer pipew.Close()

			for k := range allKeys {
				s := k.Pretty() + "\n"
				if _, err := pipew.Write([]byte(s)); err != nil {
					eptr.SetError(err)
					return
				}
			}
		}()

		res.SetOutput(eptr)
	},
}

func objectsForPaths(n *core.IpfsNode, paths []string) ([]*dag.Node, error) {
	objects := make([]*dag.Node, len(paths))
	for i, p := range paths {
		o, err := n.Resolver.ResolvePath(path.Path(p))
		if err != nil {
			return nil, err
		}
		objects[i] = o
	}
	return objects, nil
}

// ErrPassThroughReader is a reader that may return an externally set error.
type ErrPassThroughReader struct {
	R   io.ReadCloser
	err error

	sync.RWMutex
}

func (r *ErrPassThroughReader) Error() error {
	r.RLock()
	defer r.RUnlock()
	return r.err
}

func (r *ErrPassThroughReader) SetError(err error) {
	r.Lock()
	r.err = err
	r.Unlock()
}

func (r *ErrPassThroughReader) Read(buf []byte) (int, error) {
	err := r.Error()
	if err != nil {
		return 0, err
	}

	return r.R.Read(buf)
}

func (r *ErrPassThroughReader) Close() error {
	err1 := r.R.Close()
	err2 := r.Error()
	if err2 != nil {
		return err2
	}
	return err1
}

type RefWriter struct {
	W   io.Writer
	DAG dag.DAGService
	Ctx context.Context

	Unique    bool
	Recursive bool
	PrintEdge bool
	PrintFmt  string

	seen map[u.Key]struct{}
}

// WriteRefs writes refs of the given object to the underlying writer.
func (rw *RefWriter) WriteRefs(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	if rw.skip(nkey) {
		return 0, nil
	}

	count := 0
	for _, l := range n.Links {
		lk := u.Key(l.Hash)

		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, l.Name); err != nil {
			return count, err
		}
		count++

		if !rw.Recursive {
			continue
		}

		child, err := l.GetNode(rw.DAG)
		if err != nil {
			return count, err
		}

		c, err := rw.WriteRefs(child)
		count += c
		if err != nil {
			return count, err
		}
	}
	return count, nil
}

// skip returns whether to skip a key
func (rw *RefWriter) skip(k u.Key) bool {
	if !rw.Unique {
		return false
	}

	if rw.seen == nil {
		rw.seen = make(map[u.Key]struct{})
	}

	_, found := rw.seen[k]
	if !found {
		rw.seen[k] = struct{}{}
	}
	return found
}

// Write one edge
func (rw *RefWriter) WriteEdge(from, to u.Key, linkname string) error {
	if rw.Ctx != nil {
		select {
		case <-rw.Ctx.Done(): // just in case.
			return rw.Ctx.Err()
		default:
		}
	}

	var s string
	switch {
	case rw.PrintFmt != "":
		s = rw.PrintFmt
		s = strings.Replace(s, "<src>", from.Pretty(), -1)
		s = strings.Replace(s, "<dst>", to.Pretty(), -1)
		s = strings.Replace(s, "<linkname>", linkname, -1)
	case rw.PrintEdge:
		s = from.Pretty() + " -> " + to.Pretty()
	default:
		s += to.Pretty()
	}
	s += "\n"

	if _, err := rw.W.Write([]byte(s)); err != nil {
		return err
	}
	return nil
}
<commit_msg>teach ipfs refs -r how to use GetBlocks<commit_after>package commands

import (
	"bytes"
	"io"
	"strings"
	"sync"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"

	cmds "github.com/jbenet/go-ipfs/commands"
	"github.com/jbenet/go-ipfs/core"
	dag "github.com/jbenet/go-ipfs/merkledag"
	path "github.com/jbenet/go-ipfs/path"
	u "github.com/jbenet/go-ipfs/util"
)

// KeyList is a general type for outputting lists of keys
type KeyList struct {
	Keys []u.Key
}

// KeyListTextMarshaler outputs a KeyList as plaintext, one key per line
func KeyListTextMarshaler(res cmds.Response) (io.Reader, error) {
	output := res.Output().(*KeyList)
	var buf bytes.Buffer
	for _, key := range output.Keys {
		buf.WriteString(key.B58String() + "\n")
	}
	return &buf, nil
}

var RefsCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists links (references) from an object",
		ShortDescription: `
Retrieves the object named by <ipfs-path> and displays the link
hashes it contains, with the following format:

  <link base58 hash>

Note: list all refs recursively with -r.
`,
	},
	Subcommands: map[string]*cmds.Command{
		"local": RefsLocalCmd,
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("ipfs-path", true, true, "Path to the object(s) to list refs from").EnableStdin(),
	},
	Options: []cmds.Option{
		cmds.StringOption("format", "Emit edges with given format. tokens: <src> <dst> <linkname>"),
		cmds.BoolOption("edges", "e", "Emit edge format: `<from> -> <to>`"),
		cmds.BoolOption("unique", "u", "Omit duplicate refs from output"),
		cmds.BoolOption("recursive", "r", "Recursively list links of child nodes"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context().Context
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		unique, _, err := req.Option("unique").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		recursive, _, err := req.Option("recursive").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		edges, _, err := req.Option("edges").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		format, _, err := req.Option("format").String()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		objs, err := objectsForPaths(n, req.Arguments())
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		piper, pipew := io.Pipe()
		eptr := &ErrPassThroughReader{R: piper}

		go func() {
			defer pipew.Close()

			rw := RefWriter{
				W:         pipew,
				DAG:       n.DAG,
				Ctx:       ctx,
				Unique:    unique,
				PrintEdge: edges,
				PrintFmt:  format,
				Recursive: recursive,
			}

			for _, o := range objs {
				if _, err := rw.WriteRefs(o); err != nil {
					eptr.SetError(err)
					return
				}
			}
		}()

		res.SetOutput(eptr)
	},
}

var RefsLocalCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists all local references",
		ShortDescription: `
Displays the hashes of all local objects.
`,
	},

	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context().Context
		n, err := req.Context().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		// todo: make async
		allKeys, err := n.Blockstore.AllKeysChan(ctx)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		piper, pipew := io.Pipe()
		eptr := &ErrPassThroughReader{R: piper}

		go func() {
			defer pipew.Close()

			for k := range allKeys {
				s := k.Pretty() + "\n"
				if _, err := pipew.Write([]byte(s)); err != nil {
					eptr.SetError(err)
					return
				}
			}
		}()

		res.SetOutput(eptr)
	},
}

func objectsForPaths(n *core.IpfsNode, paths []string) ([]*dag.Node, error) {
	objects := make([]*dag.Node, len(paths))
	for i, p := range paths {
		o, err := n.Resolver.ResolvePath(path.Path(p))
		if err != nil {
			return nil, err
		}
		objects[i] = o
	}
	return objects, nil
}

// ErrPassThroughReader is a reader that may return an externally set error.
type ErrPassThroughReader struct {
	R   io.ReadCloser
	err error

	sync.RWMutex
}

func (r *ErrPassThroughReader) Error() error {
	r.RLock()
	defer r.RUnlock()
	return r.err
}

func (r *ErrPassThroughReader) SetError(err error) {
	r.Lock()
	r.err = err
	r.Unlock()
}

func (r *ErrPassThroughReader) Read(buf []byte) (int, error) {
	err := r.Error()
	if err != nil {
		return 0, err
	}

	return r.R.Read(buf)
}

func (r *ErrPassThroughReader) Close() error {
	err1 := r.R.Close()
	err2 := r.Error()
	if err2 != nil {
		return err2
	}
	return err1
}

type RefWriter struct {
	W   io.Writer
	DAG dag.DAGService
	Ctx context.Context

	Unique    bool
	Recursive bool
	PrintEdge bool
	PrintFmt  string

	seen map[u.Key]struct{}
}

// WriteRefs writes refs of the given object to the underlying writer.
func (rw *RefWriter) WriteRefs(n *dag.Node) (int, error) {
	if rw.Recursive {
		return rw.writeRefsRecursive(n)
	} else {
		return rw.writeRefsSingle(n)
	}
}

func (rw *RefWriter) writeRefsRecursive(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	if rw.skip(nkey) {
		return 0, nil
	}

	var count int
	for i, ng := range rw.DAG.GetDAG(rw.Ctx, n) {
		lk := u.Key(n.Links[i].Hash)
		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, n.Links[i].Name); err != nil {
			return count, err
		}

		nd, err := ng.Get()
		if err != nil {
			return count, err
		}

		c, err := rw.writeRefsRecursive(nd)
		count += c
		if err != nil {
			return count, err
		}
	}
	return count, nil
}

func (rw *RefWriter) writeRefsSingle(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	if rw.skip(nkey) {
		return 0, nil
	}

	count := 0
	for _, l := range n.Links {
		lk := u.Key(l.Hash)

		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, l.Name); err != nil {
			return count, err
		}
		count++
	}
	return count, nil
}

// skip returns whether to skip a key
func (rw *RefWriter) skip(k u.Key) bool {
	if !rw.Unique {
		return false
	}

	if rw.seen == nil {
		rw.seen = make(map[u.Key]struct{})
	}

	_, found := rw.seen[k]
	if !found {
		rw.seen[k] = struct{}{}
	}
	return found
}

// Write one edge
func (rw *RefWriter) WriteEdge(from, to u.Key, linkname string) error {
	if rw.Ctx != nil {
		select {
		case <-rw.Ctx.Done(): // just in case.
			return rw.Ctx.Err()
		default:
		}
	}

	var s string
	switch {
	case rw.PrintFmt != "":
		s = rw.PrintFmt
		s = strings.Replace(s, "<src>", from.Pretty(), -1)
		s = strings.Replace(s, "<dst>", to.Pretty(), -1)
		s = strings.Replace(s, "<linkname>", linkname, -1)
	case rw.PrintEdge:
		s = from.Pretty() + " -> " + to.Pretty()
	default:
		s += to.Pretty()
	}
	s += "\n"

	if _, err := rw.W.Write([]byte(s)); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package commands

import (
	"io"
	"strings"

	cmds "github.com/jbenet/go-ipfs/commands"
	u "github.com/jbenet/go-ipfs/util"
)

var log = u.Logger("core/commands")

type TestOutput struct {
	Foo string
	Bar int
}

var Root = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "global p2p merkle-dag filesystem",
		Synopsis: `
ipfs [<flags>] <command> [<arg>] ...
`,
		ShortDescription: `
Basic commands:

    init          Initialize ipfs local configuration
    add <path>    Add an object to ipfs
    cat <ref>     Show ipfs object data
    ls <ref>      List links from an object
    refs <ref>    List hashes of links from an object

Tool commands:

    config        Manage configuration
    update        Download and apply go-ipfs updates
    version       Show ipfs version information
    commands      List all available commands
    id            Show info about ipfs peers
    pin           Pin objects to local storage
    name          Publish or resolve IPNS names
    log           Change the logging level

Network commands:

    swarm         Manage connections to the p2p network
    bootstrap     Add or remove bootstrap peers
    ping          Measure the latency of a connection

Advanced Commands:

    daemon        Start a long-running daemon process
    mount         Mount an ipfs read-only mountpoint
    diag          Print diagnostics

Plumbing commands:

    block         Interact with raw blocks in the datastore
    object        Interact with raw dag nodes

Use 'ipfs <command> --help' to learn more about each command.
`,
	},
	Options: []cmds.Option{
		cmds.StringOption("config", "c", "Path to the configuration file to use"),
		cmds.BoolOption("debug", "D", "Operate in debug mode"),
		cmds.BoolOption("help", "Show the full command help text"),
		cmds.BoolOption("h", "Show a short version of the command help text"),
		cmds.BoolOption("local", "L", "Run the command locally, instead of using the daemon"),
	},
}

// commandsDaemonCmd is the "ipfs commands" command for daemon
var CommandsDaemonCmd = CommandsCmd(Root)

var rootSubcommands = map[string]*cmds.Command{
	"add":       AddCmd,
	"block":     BlockCmd,
	"bootstrap": BootstrapCmd,
	"cat":       CatCmd,
	"commands":  CommandsDaemonCmd,
	"config":    ConfigCmd,
	"diag":      DiagCmd,
	"id":        IDCmd,
	"log":       LogCmd,
	"ls":        LsCmd,
	"mount":     MountCmd,
	"name":      NameCmd,
	"object":    ObjectCmd,
	"pin":       PinCmd,
	"ping":      PingCmd,
	"refs":      RefsCmd,
	"repo":      RepoCmd,
	"swarm":     SwarmCmd,
	"update":    UpdateCmd,
	"version":   VersionCmd,
}

func init() {
	Root.Subcommands = rootSubcommands
}

type MessageOutput struct {
	Message string
}

func MessageTextMarshaler(res cmds.Response) (io.Reader, error) {
	return strings.NewReader(res.Output().(*MessageOutput).Message), nil
}
<commit_msg>core/commands: Rearranged command list sections<commit_after>package commands

import (
	"io"
	"strings"

	cmds "github.com/jbenet/go-ipfs/commands"
	u "github.com/jbenet/go-ipfs/util"
)

var log = u.Logger("core/commands")

type TestOutput struct {
	Foo string
	Bar int
}

var Root = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "global p2p merkle-dag filesystem",
		Synopsis: `
ipfs [<flags>] <command> [<arg>] ...
`,
		ShortDescription: `
Basic commands:

    init          Initialize ipfs local configuration
    add <path>    Add an object to ipfs
    cat <ref>     Show ipfs object data
    ls <ref>      List links from an object
    refs <ref>    List hashes of links from an object

Tool commands:

    config        Manage configuration
    update        Download and apply go-ipfs updates
    version       Show ipfs version information
    commands      List all available commands
    id            Show info about ipfs peers
    pin           Pin objects to local storage
    name          Publish or resolve IPNS names
    log           Change the logging level

Advanced Commands:

    daemon        Start a long-running daemon process
    mount         Mount an ipfs read-only mountpoint
    diag          Print diagnostics

Network commands:

    swarm         Manage connections to the p2p network
    bootstrap     Add or remove bootstrap peers
    ping          Measure the latency of a connection

Plumbing commands:

    block         Interact with raw blocks in the datastore
    object        Interact with raw dag nodes

Use 'ipfs <command> --help' to learn more about each command.
`,
	},
	Options: []cmds.Option{
		cmds.StringOption("config", "c", "Path to the configuration file to use"),
		cmds.BoolOption("debug", "D", "Operate in debug mode"),
		cmds.BoolOption("help", "Show the full command help text"),
		cmds.BoolOption("h", "Show a short version of the command help text"),
		cmds.BoolOption("local", "L", "Run the command locally, instead of using the daemon"),
	},
}

// commandsDaemonCmd is the "ipfs commands" command for daemon
var CommandsDaemonCmd = CommandsCmd(Root)

var rootSubcommands = map[string]*cmds.Command{
	"add":       AddCmd,
	"block":     BlockCmd,
	"bootstrap": BootstrapCmd,
	"cat":       CatCmd,
	"commands":  CommandsDaemonCmd,
	"config":    ConfigCmd,
	"diag":      DiagCmd,
	"id":        IDCmd,
	"log":       LogCmd,
	"ls":        LsCmd,
	"mount":     MountCmd,
	"name":      NameCmd,
	"object":    ObjectCmd,
	"pin":       PinCmd,
	"ping":      PingCmd,
	"refs":      RefsCmd,
	"repo":      RepoCmd,
	"swarm":     SwarmCmd,
	"update":    UpdateCmd,
	"version":   VersionCmd,
}

func init() {
	Root.Subcommands = rootSubcommands
}

type MessageOutput struct {
	Message string
}

func MessageTextMarshaler(res cmds.Response) (io.Reader, error) {
	return strings.NewReader(res.Output().(*MessageOutput).Message), nil
}
<|endoftext|>


<commit_before>// Copyright 2013 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 api

import (
	"fmt"
	"github.com/bmizerany/pat"
	"github.com/globocom/config"
	"github.com/globocom/tsuru/app"
	"github.com/globocom/tsuru/db"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"net"
	"net/http"
)

func fatal(err error) {
	log.Fatal(err.Error())
}

// RunServer starts Tsuru API server. The dry parameter indicates whether the
// server should run in dry mode, not starting the HTTP listener (for testing
// purposes).
func RunServer(dry bool) {
	log.Init()
	connString, err := config.GetString("database:url")
	if err != nil {
		connString = db.DefaultDatabaseURL
	}
	dbName, err := config.GetString("database:name")
	if err != nil {
		dbName = db.DefaultDatabaseName
	}
	fmt.Printf("Using the database %q from the server %q.\n\n", dbName, connString)

	m := pat.New()

	m.Get("/schema/app", authorizationRequiredHandler(appSchema))
	m.Get("/schema/service", authorizationRequiredHandler(serviceSchema))
	m.Get("/schema/services", authorizationRequiredHandler(servicesSchema))

	m.Get("/services/instances", authorizationRequiredHandler(serviceInstances))
	m.Get("/services/instances/:name", authorizationRequiredHandler(serviceInstance))
	m.Del("/services/instances/:name", authorizationRequiredHandler(removeServiceInstance))
	m.Post("/services/instances", authorizationRequiredHandler(createServiceInstance))
	m.Put("/services/instances/:instance/:app", authorizationRequiredHandler(bindServiceInstance))
	m.Del("/services/instances/:instance/:app", authorizationRequiredHandler(unbindServiceInstance))
	m.Get("/services/instances/:instance/status", authorizationRequiredHandler(serviceInstanceStatus))

	m.Get("/services", authorizationRequiredHandler(serviceList))
	m.Post("/services", authorizationRequiredHandler(serviceCreate))
	m.Put("/services", authorizationRequiredHandler(serviceUpdate))
	m.Del("/services/:name", authorizationRequiredHandler(serviceDelete))
	m.Get("/services/:name", authorizationRequiredHandler(serviceInfo))
	m.Get("/services/:name/doc", authorizationRequiredHandler(serviceDoc))
	m.Put("/services/:name/doc", authorizationRequiredHandler(serviceAddDoc))
	m.Put("/services/:service/:team", authorizationRequiredHandler(grantServiceAccess))
	m.Del("/services/:service/:team", authorizationRequiredHandler(revokeServiceAccess))

	m.Del("/apps/:app", authorizationRequiredHandler(appDelete))
	m.Get("/apps/:app", authorizationRequiredHandler(appInfo))
	m.Post("/apps/:app/cname", authorizationRequiredHandler(setCName))
	m.Del("/apps/:app/cname", authorizationRequiredHandler(unsetCName))
	m.Post("/apps/:app/run", authorizationRequiredHandler(runCommand))
	m.Get("/apps/:app/restart", authorizationRequiredHandler(restart))
	m.Get("/apps/:app/start", authorizationRequiredHandler(start))
	m.Get("/apps/:app/env", authorizationRequiredHandler(getEnv))
	m.Post("/apps/:app/env", authorizationRequiredHandler(setEnv))
	m.Del("/apps/:app/env", authorizationRequiredHandler(unsetEnv))
	m.Get("/apps", authorizationRequiredHandler(appList))
	m.Post("/apps", authorizationRequiredHandler(createApp))
	m.Put("/apps/:app/units", authorizationRequiredHandler(addUnits))
	m.Del("/apps/:app/units", authorizationRequiredHandler(removeUnits))
	m.Put("/apps/:app/:team", authorizationRequiredHandler(grantAppAccess))
	m.Del("/apps/:app/:team", authorizationRequiredHandler(revokeAppAccess))
	m.Get("/apps/:app/log", authorizationRequiredHandler(appLog))
	m.Post("/apps/:app/log", authorizationRequiredHandler(addLog))

	m.Get("/deploys", adminRequiredHandler(deploysList))

	m.Get("/platforms", authorizationRequiredHandler(platformList))

	// These handlers don't use :app on purpose. Using :app means that only
	// the token generate for the given app is valid, but these handlers
	// use a token generated for Gandalf.
	m.Get("/apps/:appname/available", authorizationRequiredHandler(appIsAvailable))
	m.Post("/apps/:appname/repository/clone", authorizationRequiredHandler(deploy))
	m.Post("/apps/:appname/deploy", authorizationRequiredHandler(deploy))

	if registrationEnabled, _ := config.GetBool("auth:user-registration"); registrationEnabled {
		m.Post("/users", Handler(createUser))
	}

	m.Post("/users/:email/password", Handler(resetPassword))
	m.Post("/users/:email/tokens", Handler(login))
	m.Del("/users/tokens", authorizationRequiredHandler(logout))
	m.Put("/users/password", authorizationRequiredHandler(changePassword))
	m.Del("/users", authorizationRequiredHandler(removeUser))
	m.Get("/users/:email/keys", authorizationRequiredHandler(listKeys))
	m.Post("/users/keys", authorizationRequiredHandler(addKeyToUser))
	m.Del("/users/keys", authorizationRequiredHandler(removeKeyFromUser))

	m.Post("/tokens", adminRequiredHandler(generateAppToken))

	m.Del("/logs", adminRequiredHandler(logRemove))

	m.Get("/teams", authorizationRequiredHandler(teamList))
	m.Post("/teams", authorizationRequiredHandler(createTeam))
	m.Get("/teams/:name", authorizationRequiredHandler(getTeam))
	m.Del("/teams/:name", authorizationRequiredHandler(removeTeam))
	m.Put("/teams/:team/:user", authorizationRequiredHandler(addUserToTeam))
	m.Del("/teams/:team/:user", authorizationRequiredHandler(removeUserFromTeam))

	m.Get("/healers", authorizationRequiredHandler(healers))
	m.Get("/healers/:healer", authorizationRequiredHandler(healer))

	m.Put("/swap", authorizationRequiredHandler(swap))

	m.Get("/healthcheck/", http.HandlerFunc(healthcheck))

	if !dry {
		provisioner, err := getProvisioner()
		if err != nil {
			fmt.Printf("Warning: configuration didn't declare a provisioner, using default provisioner.\n")
		}
		app.Provisioner, err = provision.Get(provisioner)
		if err != nil {
			fatal(err)
		}
		fmt.Printf("Using %q provisioner.\n\n", provisioner)

		listen, err := config.GetString("listen")
		if err != nil {
			fatal(err)
		}
		tls, _ := config.GetBool("use-tls")
		if tls {
			certFile, err := config.GetString("tls:cert-file")
			if err != nil {
				fatal(err)
			}
			keyFile, err := config.GetString("tls:key-file")
			if err != nil {
				fatal(err)
			}
			fmt.Printf("tsuru HTTP/TLS server listening at %s...\n", listen)
			fatal(http.ListenAndServeTLS(listen, certFile, keyFile, m))
		} else {
			listener, err := net.Listen("tcp", listen)
			if err != nil {
				fatal(err)
			}
			fmt.Printf("tsuru HTTP server listening at %s...\n", listen)
			http.Handle("/", m)
			fatal(http.Serve(listener, nil))
		}
	}
}
<commit_msg>api: added route for service plans.<commit_after>// Copyright 2013 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 api

import (
	"fmt"
	"github.com/bmizerany/pat"
	"github.com/globocom/config"
	"github.com/globocom/tsuru/app"
	"github.com/globocom/tsuru/db"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"net"
	"net/http"
)

func fatal(err error) {
	log.Fatal(err.Error())
}

// RunServer starts Tsuru API server. The dry parameter indicates whether the
// server should run in dry mode, not starting the HTTP listener (for testing
// purposes).
func RunServer(dry bool) {
	log.Init()
	connString, err := config.GetString("database:url")
	if err != nil {
		connString = db.DefaultDatabaseURL
	}
	dbName, err := config.GetString("database:name")
	if err != nil {
		dbName = db.DefaultDatabaseName
	}
	fmt.Printf("Using the database %q from the server %q.\n\n", dbName, connString)

	m := pat.New()

	m.Get("/schema/app", authorizationRequiredHandler(appSchema))
	m.Get("/schema/service", authorizationRequiredHandler(serviceSchema))
	m.Get("/schema/services", authorizationRequiredHandler(servicesSchema))

	m.Get("/services/instances", authorizationRequiredHandler(serviceInstances))
	m.Get("/services/instances/:name", authorizationRequiredHandler(serviceInstance))
	m.Del("/services/instances/:name", authorizationRequiredHandler(removeServiceInstance))
	m.Post("/services/instances", authorizationRequiredHandler(createServiceInstance))
	m.Put("/services/instances/:instance/:app", authorizationRequiredHandler(bindServiceInstance))
	m.Del("/services/instances/:instance/:app", authorizationRequiredHandler(unbindServiceInstance))
	m.Get("/services/instances/:instance/status", authorizationRequiredHandler(serviceInstanceStatus))

	m.Get("/services", authorizationRequiredHandler(serviceList))
	m.Post("/services", authorizationRequiredHandler(serviceCreate))
	m.Put("/services", authorizationRequiredHandler(serviceUpdate))
	m.Del("/services/:name", authorizationRequiredHandler(serviceDelete))
	m.Get("/services/:name", authorizationRequiredHandler(serviceInfo))
	m.Get("/services/:name/plans", authorizationRequiredHandler(servicePlans))
	m.Get("/services/:name/doc", authorizationRequiredHandler(serviceDoc))
	m.Put("/services/:name/doc", authorizationRequiredHandler(serviceAddDoc))
	m.Put("/services/:service/:team", authorizationRequiredHandler(grantServiceAccess))
	m.Del("/services/:service/:team", authorizationRequiredHandler(revokeServiceAccess))

	m.Del("/apps/:app", authorizationRequiredHandler(appDelete))
	m.Get("/apps/:app", authorizationRequiredHandler(appInfo))
	m.Post("/apps/:app/cname", authorizationRequiredHandler(setCName))
	m.Del("/apps/:app/cname", authorizationRequiredHandler(unsetCName))
	m.Post("/apps/:app/run", authorizationRequiredHandler(runCommand))
	m.Get("/apps/:app/restart", authorizationRequiredHandler(restart))
	m.Get("/apps/:app/start", authorizationRequiredHandler(start))
	m.Get("/apps/:app/env", authorizationRequiredHandler(getEnv))
	m.Post("/apps/:app/env", authorizationRequiredHandler(setEnv))
	m.Del("/apps/:app/env", authorizationRequiredHandler(unsetEnv))
	m.Get("/apps", authorizationRequiredHandler(appList))
	m.Post("/apps", authorizationRequiredHandler(createApp))
	m.Put("/apps/:app/units", authorizationRequiredHandler(addUnits))
	m.Del("/apps/:app/units", authorizationRequiredHandler(removeUnits))
	m.Put("/apps/:app/:team", authorizationRequiredHandler(grantAppAccess))
	m.Del("/apps/:app/:team", authorizationRequiredHandler(revokeAppAccess))
	m.Get("/apps/:app/log", authorizationRequiredHandler(appLog))
	m.Post("/apps/:app/log", authorizationRequiredHandler(addLog))

	m.Get("/deploys", adminRequiredHandler(deploysList))

	m.Get("/platforms", authorizationRequiredHandler(platformList))

	// These handlers don't use :app on purpose. Using :app means that only
	// the token generate for the given app is valid, but these handlers
	// use a token generated for Gandalf.
	m.Get("/apps/:appname/available", authorizationRequiredHandler(appIsAvailable))
	m.Post("/apps/:appname/repository/clone", authorizationRequiredHandler(deploy))
	m.Post("/apps/:appname/deploy", authorizationRequiredHandler(deploy))

	if registrationEnabled, _ := config.GetBool("auth:user-registration"); registrationEnabled {
		m.Post("/users", Handler(createUser))
	}

	m.Post("/users/:email/password", Handler(resetPassword))
	m.Post("/users/:email/tokens", Handler(login))
	m.Del("/users/tokens", authorizationRequiredHandler(logout))
	m.Put("/users/password", authorizationRequiredHandler(changePassword))
	m.Del("/users", authorizationRequiredHandler(removeUser))
	m.Get("/users/:email/keys", authorizationRequiredHandler(listKeys))
	m.Post("/users/keys", authorizationRequiredHandler(addKeyToUser))
	m.Del("/users/keys", authorizationRequiredHandler(removeKeyFromUser))

	m.Post("/tokens", adminRequiredHandler(generateAppToken))

	m.Del("/logs", adminRequiredHandler(logRemove))

	m.Get("/teams", authorizationRequiredHandler(teamList))
	m.Post("/teams", authorizationRequiredHandler(createTeam))
	m.Get("/teams/:name", authorizationRequiredHandler(getTeam))
	m.Del("/teams/:name", authorizationRequiredHandler(removeTeam))
	m.Put("/teams/:team/:user", authorizationRequiredHandler(addUserToTeam))
	m.Del("/teams/:team/:user", authorizationRequiredHandler(removeUserFromTeam))

	m.Get("/healers", authorizationRequiredHandler(healers))
	m.Get("/healers/:healer", authorizationRequiredHandler(healer))

	m.Put("/swap", authorizationRequiredHandler(swap))

	m.Get("/healthcheck/", http.HandlerFunc(healthcheck))

	if !dry {
		provisioner, err := getProvisioner()
		if err != nil {
			fmt.Printf("Warning: configuration didn't declare a provisioner, using default provisioner.\n")
		}
		app.Provisioner, err = provision.Get(provisioner)
		if err != nil {
			fatal(err)
		}
		fmt.Printf("Using %q provisioner.\n\n", provisioner)

		listen, err := config.GetString("listen")
		if err != nil {
			fatal(err)
		}
		tls, _ := config.GetBool("use-tls")
		if tls {
			certFile, err := config.GetString("tls:cert-file")
			if err != nil {
				fatal(err)
			}
			keyFile, err := config.GetString("tls:key-file")
			if err != nil {
				fatal(err)
			}
			fmt.Printf("tsuru HTTP/TLS server listening at %s...\n", listen)
			fatal(http.ListenAndServeTLS(listen, certFile, keyFile, m))
		} else {
			listener, err := net.Listen("tcp", listen)
			if err != nil {
				fatal(err)
			}
			fmt.Printf("tsuru HTTP server listening at %s...\n", listen)
			http.Handle("/", m)
			fatal(http.Serve(listener, nil))
		}
	}
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package api

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/gorilla/handlers"

	"github.com/rs/cors"

	"github.com/ava-labs/avalanchego/api/auth"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/engine/common"
	"github.com/ava-labs/avalanchego/utils/logging"
)

const (
	baseURL               = "/ext"
	serverShutdownTimeout = 10 * time.Second
)

var (
	errUnknownLockOption = errors.New("invalid lock options")
)

// Server maintains the HTTP router
type Server struct {
	// log this server writes to
	log logging.Logger
	// generates new logs for chains to write to
	factory logging.Factory
	// Maps endpoints to handlers
	router *router
	// points the the router handlers
	handler http.Handler
	// Listens for HTTP traffic on this address
	listenAddress string
	// Handles authorization. Must be non-nil after initialization, even if
	// token authorization is off.
	auth *auth.Auth

	// http server
	srv *http.Server
}

// Initialize creates the API server at the provided host and port
func (s *Server) Initialize(
	log logging.Logger,
	factory logging.Factory,
	host string,
	port uint16,
	authEnabled bool,
	authPassword string,
	allowedOrigins []string,
) error {
	s.log = log
	s.factory = factory
	s.listenAddress = fmt.Sprintf("%s:%d", host, port)
	s.router = newRouter()

	a, err := auth.New(authEnabled, authPassword)
	if err != nil {
		return err
	}
	s.auth = a

	s.log.Info("API created with allowed origins: %v", allowedOrigins)
	corsWrapper := cors.New(cors.Options{
		AllowedOrigins: allowedOrigins,
	})
	corsHandler := corsWrapper.Handler(s.router)
	s.handler = s.auth.WrapHandler(corsHandler)

	if !authEnabled {
		return nil
	}

	// only create auth service if token authorization is required
	s.log.Info("API authorization is enabled. Auth tokens must be passed in the header of API requests, except requests to the auth service.")
	authService := auth.NewService(s.log, s.auth)
	return s.AddRoute(authService, &sync.RWMutex{}, auth.Endpoint, "", s.log)

}

// Dispatch starts the API server
func (s *Server) Dispatch() error {
	listener, err := net.Listen("tcp", s.listenAddress)
	if err != nil {
		return err
	}
	s.log.Info("HTTP API server listening on %q", s.listenAddress)
	s.srv = &http.Server{Handler: s.handler}
	return s.srv.Serve(listener)
}

// DispatchTLS starts the API server with the provided TLS certificate
func (s *Server) DispatchTLS(certFile, keyFile string) error {
	listener, err := net.Listen("tcp", s.listenAddress)
	if err != nil {
		return err
	}
	s.log.Info("HTTPS API server listening on %q", s.listenAddress)
	return http.ServeTLS(listener, s.handler, certFile, keyFile)
}

// RegisterChain registers the API endpoints associated with this chain That is,
// add <route, handler> pairs to server so that http calls can be made to the vm
func (s *Server) RegisterChain(chainName string, ctx *snow.Context, vmIntf interface{}) {
	vm, ok := vmIntf.(common.VM)
	if !ok {
		return
	}

	ctx.Lock.Lock()
	handlers, err := vm.CreateHandlers()
	ctx.Lock.Unlock()
	if err != nil {
		s.log.Error("Failed to create %s handlers: %s", chainName, err)
		return
	}

	httpLogger, err := s.factory.MakeChain(chainName, "http")
	if err != nil {
		s.log.Error("Failed to create new http logger: %s", err)
		return
	}

	s.log.Verbo("About to add API endpoints for chain with ID %s", ctx.ChainID)
	// all subroutes to a chain begin with "bc/<the chain's ID>"
	defaultEndpoint := "bc/" + ctx.ChainID.String()

	// Register each endpoint
	for extension, service := range handlers {
		// Validate that the route being added is valid
		// e.g. "/foo" and "" are ok but "\n" is not
		_, err := url.ParseRequestURI(extension)
		if extension != "" && err != nil {
			s.log.Error("could not add route to chain's API handler because route is malformed: %s", err)
			continue
		}
		if err := s.AddChainRoute(service, ctx, defaultEndpoint, extension, httpLogger); err != nil {
			s.log.Error("error adding route: %s", err)
		}
	}
}

// AddChainRoute registers a route to a chain's handler
func (s *Server) AddChainRoute(handler *common.HTTPHandler, ctx *snow.Context, base, endpoint string, loggingWriter io.Writer) error {
	url := fmt.Sprintf("%s/%s", baseURL, base)
	s.log.Info("adding route %s%s", url, endpoint)
	// Apply logging middleware
	h := handlers.CombinedLoggingHandler(loggingWriter, handler.Handler)
	// Apply middleware to grab/release chain's lock before/after calling API method
	h, err := lockMiddleware(h, handler.LockOptions, &ctx.Lock)
	if err != nil {
		return err
	}
	// Apply middleware to reject calls to the handler before the chain finishes bootstrapping
	h = rejectMiddleware(h, ctx)
	return s.router.AddRouter(url, endpoint, h)
}

// AddRoute registers a route to a handler.
func (s *Server) AddRoute(handler *common.HTTPHandler, lock *sync.RWMutex, base, endpoint string, loggingWriter io.Writer) error {
	url := fmt.Sprintf("%s/%s", baseURL, base)
	s.log.Info("adding route %s%s", url, endpoint)
	// Apply logging middleware
	h := handlers.CombinedLoggingHandler(loggingWriter, handler.Handler)
	// Apply middleware to grab/release chain's lock before/after calling API method
	h, err := lockMiddleware(h, handler.LockOptions, lock)
	if err != nil {
		return err
	}
	return s.router.AddRouter(url, endpoint, h)
}

// Wraps a handler by grabbing and releasing a lock before calling the handler.
func lockMiddleware(handler http.Handler, lockOption common.LockOption, lock *sync.RWMutex) (http.Handler, error) {
	switch lockOption {
	case common.WriteLock:
		return middlewareHandler{
			before:  lock.Lock,
			after:   lock.Unlock,
			handler: handler,
		}, nil
	case common.ReadLock:
		return middlewareHandler{
			before:  lock.RLock,
			after:   lock.RUnlock,
			handler: handler,
		}, nil
	case common.NoLock:
		return handler, nil
	default:
		return nil, errUnknownLockOption
	}
}

// Reject middleware wraps a handler. If the chain that the context describes is
// not done bootstrapping, writes back an error.
func rejectMiddleware(handler http.Handler, ctx *snow.Context) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // If chain isn't done bootstrapping, ignore API calls
		if !ctx.IsBootstrapped() {
			w.WriteHeader(http.StatusServiceUnavailable)
			// Doesn't matter if there's an error while writing. They'll get the StatusServiceUnavailable code.
			_, _ = w.Write([]byte("API call rejected because chain is not done bootstrapping"))
		} else {
			handler.ServeHTTP(w, r)
		}
	})
}

// AddAliases registers aliases to the server
func (s *Server) AddAliases(endpoint string, aliases ...string) error {
	url := fmt.Sprintf("%s/%s", baseURL, endpoint)
	endpoints := make([]string, len(aliases))
	for i, alias := range aliases {
		endpoints[i] = fmt.Sprintf("%s/%s", baseURL, alias)
	}
	return s.router.AddAlias(url, endpoints...)
}

// AddAliasesWithReadLock registers aliases to the server assuming the http read
// lock is currently held.
func (s *Server) AddAliasesWithReadLock(endpoint string, aliases ...string) error {
	// This is safe, as the read lock doesn't actually need to be held once the
	// http handler is called. However, it is unlocked later, so this function
	// must end with the lock held.
	s.router.lock.RUnlock()
	defer s.router.lock.RLock()

	return s.AddAliases(endpoint, aliases...)
}

// Call ...
func (s *Server) Call(
	writer http.ResponseWriter,
	method,
	base,
	endpoint string,
	body io.Reader,
	headers map[string]string,
) error {
	url := fmt.Sprintf("%s/vm/%s", baseURL, base)

	handler, err := s.router.GetHandler(url, endpoint)
	if err != nil {
		return err
	}
	req, err := http.NewRequest("POST", "*", body)
	if err != nil {
		return err
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}

	handler.ServeHTTP(writer, req)

	return nil
}

// Shutdown this server
func (s *Server) Shutdown() error {
	if s.srv == nil {
		return nil
	}
	ctx, cancel := context.WithTimeout(context.Background(), serverShutdownTimeout)
	defer cancel()
	return s.srv.Shutdown(ctx)
}
<commit_msg>allow credentials<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package api

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/gorilla/handlers"

	"github.com/rs/cors"

	"github.com/ava-labs/avalanchego/api/auth"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/engine/common"
	"github.com/ava-labs/avalanchego/utils/logging"
)

const (
	baseURL               = "/ext"
	serverShutdownTimeout = 10 * time.Second
)

var (
	errUnknownLockOption = errors.New("invalid lock options")
)

// Server maintains the HTTP router
type Server struct {
	// log this server writes to
	log logging.Logger
	// generates new logs for chains to write to
	factory logging.Factory
	// Maps endpoints to handlers
	router *router
	// points the the router handlers
	handler http.Handler
	// Listens for HTTP traffic on this address
	listenAddress string
	// Handles authorization. Must be non-nil after initialization, even if
	// token authorization is off.
	auth *auth.Auth

	// http server
	srv *http.Server
}

// Initialize creates the API server at the provided host and port
func (s *Server) Initialize(
	log logging.Logger,
	factory logging.Factory,
	host string,
	port uint16,
	authEnabled bool,
	authPassword string,
	allowedOrigins []string,
) error {
	s.log = log
	s.factory = factory
	s.listenAddress = fmt.Sprintf("%s:%d", host, port)
	s.router = newRouter()

	a, err := auth.New(authEnabled, authPassword)
	if err != nil {
		return err
	}
	s.auth = a

	s.log.Info("API created with allowed origins: %v", allowedOrigins)
	corsWrapper := cors.New(cors.Options{
		AllowedOrigins:   allowedOrigins,
		AllowCredentials: true,
	})
	corsHandler := corsWrapper.Handler(s.router)
	s.handler = s.auth.WrapHandler(corsHandler)

	if !authEnabled {
		return nil
	}

	// only create auth service if token authorization is required
	s.log.Info("API authorization is enabled. Auth tokens must be passed in the header of API requests, except requests to the auth service.")
	authService := auth.NewService(s.log, s.auth)
	return s.AddRoute(authService, &sync.RWMutex{}, auth.Endpoint, "", s.log)

}

// Dispatch starts the API server
func (s *Server) Dispatch() error {
	listener, err := net.Listen("tcp", s.listenAddress)
	if err != nil {
		return err
	}
	s.log.Info("HTTP API server listening on %q", s.listenAddress)
	s.srv = &http.Server{Handler: s.handler}
	return s.srv.Serve(listener)
}

// DispatchTLS starts the API server with the provided TLS certificate
func (s *Server) DispatchTLS(certFile, keyFile string) error {
	listener, err := net.Listen("tcp", s.listenAddress)
	if err != nil {
		return err
	}
	s.log.Info("HTTPS API server listening on %q", s.listenAddress)
	return http.ServeTLS(listener, s.handler, certFile, keyFile)
}

// RegisterChain registers the API endpoints associated with this chain That is,
// add <route, handler> pairs to server so that http calls can be made to the vm
func (s *Server) RegisterChain(chainName string, ctx *snow.Context, vmIntf interface{}) {
	vm, ok := vmIntf.(common.VM)
	if !ok {
		return
	}

	ctx.Lock.Lock()
	handlers, err := vm.CreateHandlers()
	ctx.Lock.Unlock()
	if err != nil {
		s.log.Error("Failed to create %s handlers: %s", chainName, err)
		return
	}

	httpLogger, err := s.factory.MakeChain(chainName, "http")
	if err != nil {
		s.log.Error("Failed to create new http logger: %s", err)
		return
	}

	s.log.Verbo("About to add API endpoints for chain with ID %s", ctx.ChainID)
	// all subroutes to a chain begin with "bc/<the chain's ID>"
	defaultEndpoint := "bc/" + ctx.ChainID.String()

	// Register each endpoint
	for extension, service := range handlers {
		// Validate that the route being added is valid
		// e.g. "/foo" and "" are ok but "\n" is not
		_, err := url.ParseRequestURI(extension)
		if extension != "" && err != nil {
			s.log.Error("could not add route to chain's API handler because route is malformed: %s", err)
			continue
		}
		if err := s.AddChainRoute(service, ctx, defaultEndpoint, extension, httpLogger); err != nil {
			s.log.Error("error adding route: %s", err)
		}
	}
}

// AddChainRoute registers a route to a chain's handler
func (s *Server) AddChainRoute(handler *common.HTTPHandler, ctx *snow.Context, base, endpoint string, loggingWriter io.Writer) error {
	url := fmt.Sprintf("%s/%s", baseURL, base)
	s.log.Info("adding route %s%s", url, endpoint)
	// Apply logging middleware
	h := handlers.CombinedLoggingHandler(loggingWriter, handler.Handler)
	// Apply middleware to grab/release chain's lock before/after calling API method
	h, err := lockMiddleware(h, handler.LockOptions, &ctx.Lock)
	if err != nil {
		return err
	}
	// Apply middleware to reject calls to the handler before the chain finishes bootstrapping
	h = rejectMiddleware(h, ctx)
	return s.router.AddRouter(url, endpoint, h)
}

// AddRoute registers a route to a handler.
func (s *Server) AddRoute(handler *common.HTTPHandler, lock *sync.RWMutex, base, endpoint string, loggingWriter io.Writer) error {
	url := fmt.Sprintf("%s/%s", baseURL, base)
	s.log.Info("adding route %s%s", url, endpoint)
	// Apply logging middleware
	h := handlers.CombinedLoggingHandler(loggingWriter, handler.Handler)
	// Apply middleware to grab/release chain's lock before/after calling API method
	h, err := lockMiddleware(h, handler.LockOptions, lock)
	if err != nil {
		return err
	}
	return s.router.AddRouter(url, endpoint, h)
}

// Wraps a handler by grabbing and releasing a lock before calling the handler.
func lockMiddleware(handler http.Handler, lockOption common.LockOption, lock *sync.RWMutex) (http.Handler, error) {
	switch lockOption {
	case common.WriteLock:
		return middlewareHandler{
			before:  lock.Lock,
			after:   lock.Unlock,
			handler: handler,
		}, nil
	case common.ReadLock:
		return middlewareHandler{
			before:  lock.RLock,
			after:   lock.RUnlock,
			handler: handler,
		}, nil
	case common.NoLock:
		return handler, nil
	default:
		return nil, errUnknownLockOption
	}
}

// Reject middleware wraps a handler. If the chain that the context describes is
// not done bootstrapping, writes back an error.
func rejectMiddleware(handler http.Handler, ctx *snow.Context) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // If chain isn't done bootstrapping, ignore API calls
		if !ctx.IsBootstrapped() {
			w.WriteHeader(http.StatusServiceUnavailable)
			// Doesn't matter if there's an error while writing. They'll get the StatusServiceUnavailable code.
			_, _ = w.Write([]byte("API call rejected because chain is not done bootstrapping"))
		} else {
			handler.ServeHTTP(w, r)
		}
	})
}

// AddAliases registers aliases to the server
func (s *Server) AddAliases(endpoint string, aliases ...string) error {
	url := fmt.Sprintf("%s/%s", baseURL, endpoint)
	endpoints := make([]string, len(aliases))
	for i, alias := range aliases {
		endpoints[i] = fmt.Sprintf("%s/%s", baseURL, alias)
	}
	return s.router.AddAlias(url, endpoints...)
}

// AddAliasesWithReadLock registers aliases to the server assuming the http read
// lock is currently held.
func (s *Server) AddAliasesWithReadLock(endpoint string, aliases ...string) error {
	// This is safe, as the read lock doesn't actually need to be held once the
	// http handler is called. However, it is unlocked later, so this function
	// must end with the lock held.
	s.router.lock.RUnlock()
	defer s.router.lock.RLock()

	return s.AddAliases(endpoint, aliases...)
}

// Call ...
func (s *Server) Call(
	writer http.ResponseWriter,
	method,
	base,
	endpoint string,
	body io.Reader,
	headers map[string]string,
) error {
	url := fmt.Sprintf("%s/vm/%s", baseURL, base)

	handler, err := s.router.GetHandler(url, endpoint)
	if err != nil {
		return err
	}
	req, err := http.NewRequest("POST", "*", body)
	if err != nil {
		return err
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}

	handler.ServeHTTP(writer, req)

	return nil
}

// Shutdown this server
func (s *Server) Shutdown() error {
	if s.srv == nil {
		return nil
	}
	ctx, cancel := context.WithTimeout(context.Background(), serverShutdownTimeout)
	defer cancel()
	return s.srv.Shutdown(ctx)
}
<|endoftext|>


<commit_before>// Copyright © 2017 Ricardo Aravena <raravena@branch.io>
//
// 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 exec

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"io/ioutil"
	"net"
	"os"
	"time"
)

type SignerContainer struct {
	signers []ssh.Signer
}

func makeSigner(keyname string) (signer ssh.Signer, err error) {
	fp, err := os.Open(keyname)
	if err != nil {
		return
	}
	defer fp.Close()

	buf, _ := ioutil.ReadAll(fp)
	signer, _ = ssh.ParsePrivateKey(buf)
	return
}

func makeKeyring(key string, useAgent bool) ssh.AuthMethod {
	signers := []ssh.Signer{}

	if useAgent == true {
		aConn, _ := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
		sshAgent := agent.NewClient(aConn)
		aSigners, _ := sshAgent.Signers()
		for _, signer := range aSigners {
			signers = append(signers, signer)
		}
	}

	keys := []string{
		key,
		os.Getenv("HOME") + "/.ssh/id_dsa"}

	for _, keyname := range keys {
		signer, err := makeSigner(keyname)
		if err == nil {
			signers = append(signers, signer)
		}
	}
	return ssh.PublicKeys(signers...)
}

func executeCmd(cmd, hostname string, config *ssh.ClientConfig) string {
	conn, err := ssh.Dial("tcp", hostname+":22", config)

	if err != nil {
		fmt.Println(err)
		return ""
	}

	session, _ := conn.NewSession()
	defer session.Close()

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Run(cmd)

	return hostname + ":\n" + stdoutBuf.String()
}

// Run the ssh command
func Run(machines []string, cmd string, user string, key string, useAgent bool) {
	// in 5 seconds the message will come to timeout channel
	timeout := time.After(5 * time.Second)
	results := make(chan string, 10)

	config := &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{makeKeyring(key, useAgent)},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	for _, m := range machines {
		go func(hostname string) {
			results <- executeCmd(cmd, hostname, config)
			// we’ll write results into the buffered channel of strings
		}(m)
	}

	for i := 0; i < len(machines); i++ {
		select {
		case res := <-results:
			fmt.Print(res)
		case <-timeout:
			fmt.Println("Timed out!")
			return
		}
	}
}
<commit_msg>Cleanup the code a bit<commit_after>// Copyright © 2017 Ricardo Aravena <raravena@branch.io>
//
// 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 exec

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"io/ioutil"
	"net"
	"os"
	"time"
)

func makeSigner(keyname string) (signer ssh.Signer, err error) {
	fp, err := os.Open(keyname)
	if err != nil {
		return
	}
	defer fp.Close()

	buf, _ := ioutil.ReadAll(fp)
	signer, _ = ssh.ParsePrivateKey(buf)
	return
}

func makeKeyring(key string, useAgent bool) ssh.AuthMethod {
	signers := []ssh.Signer{}

	if useAgent == true {
		aConn, _ := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
		sshAgent := agent.NewClient(aConn)
		aSigners, _ := sshAgent.Signers()
		for _, signer := range aSigners {
			signers = append(signers, signer)
		}
	}

	keys := []string{
		key,
		os.Getenv("HOME") + "/.ssh/id_dsa"}

	for _, keyname := range keys {
		signer, err := makeSigner(keyname)
		if err == nil {
			signers = append(signers, signer)
		}
	}
	return ssh.PublicKeys(signers...)
}

func executeCmd(cmd, hostname string, config *ssh.ClientConfig) string {
	conn, err := ssh.Dial("tcp", hostname+":22", config)

	if err != nil {
		fmt.Println(err)
		return ""
	}

	session, _ := conn.NewSession()
	defer session.Close()

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Run(cmd)

	return hostname + ":\n" + stdoutBuf.String()
}

// Run the ssh command
func Run(machines []string, cmd string, user string, key string, useAgent bool) {
	// in 5 seconds the message will come to timeout channel
	timeout := time.After(5 * time.Second)
	results := make(chan string, 10)

	config := &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{makeKeyring(key, useAgent)},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	for _, m := range machines {
		go func(hostname string) {
			results <- executeCmd(cmd, hostname, config)
			// we’ll write results into the buffered channel of strings
		}(m)
	}

	for i := 0; i < len(machines); i++ {
		select {
		case res := <-results:
			fmt.Print(res)
		case <-timeout:
			fmt.Println("Timed out!")
			return
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"log"
	"os/exec"

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

type DynoState int

const (
	Stopped DynoState = iota
	Started
	Retiring
	Retired
)

type DynoInput int

const (
	Retire DynoInput = iota
	Restart
	Exited
	StayStarted
)

var ErrExecutorComplete = errors.New("Executor complete")

type Executor struct {
	args        []string
	dynoDriver  DynoDriver
	release     *Release
	processID   string
	processType string
	status      chan *ExitStatus
	complete    chan struct{}

	// simple dyno driver properties
	cmd     *exec.Cmd
	waiting chan struct{}

	// docker dyno driver properties
	container *docker.Container

	// FSM Fields
	OneShot  bool
	state    DynoState
	newInput chan DynoInput
}

func (e *Executor) Trigger(input DynoInput) {
	log.Println("triggering", input)
	select {
	case e.newInput <- input:
	case <-e.complete:
	}
}

func (e *Executor) wait() {
	if s := e.dynoDriver.Wait(e); e.status != nil {
		log.Println("Executor exits:", e.Name(), "exit code:", s.code)
		e.status <- s
	}
	e.Trigger(Exited)
}

func (e *Executor) Tick() (err error) {
	log.Println(e.Name(), "waiting for tick...", e.state)
	input := <-e.newInput
	log.Println(e.Name(), "ticking with input", input)

	start := func() error {
		log.Printf("%v: starting\n", e.Name())
		if err = e.dynoDriver.Start(e); err != nil {
			log.Printf("%v: start fails: %v", e.Name(), err)
			go e.Trigger(Restart)
			return err
		}

		log.Printf("%v: started\n", e.Name())
		e.state = Started
		go e.wait()
		return nil
	}

again:
	switch e.state {
	case Retired:
		close(e.complete)
		return ErrExecutorComplete
	case Retiring:
		if err = e.dynoDriver.Stop(e); err != nil {
			return err
		}

		e.state = Retired
		goto again
	case Stopped:
		switch input {
		case Retire:
			e.state = Retired
			goto again
		case Exited:
			if e.OneShot {
				e.state = Retired
				goto again
			}

			return start()
		case StayStarted:
			fallthrough
		case Restart:
			return start()
		default:
			panic(fmt.Sprintln("Invalid input", input))
		}
	case Started:
		switch input {
		case Retire:
			e.state = Retiring
			goto again
		case Exited:
			e.state = Stopped
			goto again
		case Restart:
			if err = e.dynoDriver.Stop(e); err != nil {
				return err
			}
			goto again
		default:
			panic(fmt.Sprintln("Invalid input", input))
		}
	default:
		panic(fmt.Sprintln("Invalid state", e.state))
	}
}

func (e *Executor) Name() string {
	return e.processType + "." + e.processID
}
<commit_msg>Fail if a one-shot can't start<commit_after>package main

import (
	"errors"
	"log"
	"os/exec"

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

type DynoState int

const (
	Stopped DynoState = iota
	Started
	Retiring
	Retired
)

type DynoInput int

const (
	Retire DynoInput = iota
	Restart
	Exited
	StayStarted
)

var ErrExecutorComplete = errors.New("Executor complete")

type Executor struct {
	args        []string
	dynoDriver  DynoDriver
	release     *Release
	processID   string
	processType string
	status      chan *ExitStatus
	complete    chan struct{}

	// simple dyno driver properties
	cmd     *exec.Cmd
	waiting chan struct{}

	// docker dyno driver properties
	container *docker.Container

	// FSM Fields
	OneShot  bool
	state    DynoState
	newInput chan DynoInput
}

func (e *Executor) Trigger(input DynoInput) {
	log.Println("triggering", input)
	select {
	case e.newInput <- input:
	case <-e.complete:
	}
}

func (e *Executor) wait() {
	if s := e.dynoDriver.Wait(e); e.status != nil {
		log.Println("Executor exits:", e.Name(), "exit code:", s.code)
		e.status <- s
	}
	e.Trigger(Exited)
}

func (e *Executor) Tick() (err error) {
	log.Println(e.Name(), "waiting for tick...", e.state)
	input := <-e.newInput
	log.Println(e.Name(), "ticking with input", input)

	start := func() error {
		log.Printf("%v: starting\n", e.Name())
		if err = e.dynoDriver.Start(e); err != nil {
			log.Printf("%v: start fails: %v", e.Name(), err)
			if e.OneShot {
				go e.Trigger(Retire)
			} else {
				go e.Trigger(Restart)
			}
			return err
		}

		log.Printf("%v: started\n", e.Name())
		e.state = Started
		go e.wait()
		return nil
	}

again:
	switch e.state {
	case Retired:
		close(e.complete)
		return ErrExecutorComplete
	case Retiring:
		if err = e.dynoDriver.Stop(e); err != nil {
			return err
		}

		e.state = Retired
		goto again
	case Stopped:
		switch input {
		case Retire:
			e.state = Retired
			goto again
		case Exited:
			if e.OneShot {
				e.state = Retired
				goto again
			}

			return start()
		case StayStarted:
			fallthrough
		case Restart:
			return start()
		default:
			panic(fmt.Sprintln("Invalid input", input))
		}
	case Started:
		switch input {
		case Retire:
			e.state = Retiring
			goto again
		case Exited:
			e.state = Stopped
			goto again
		case Restart:
			if err = e.dynoDriver.Stop(e); err != nil {
				return err
			}
			goto again
		default:
			panic(fmt.Sprintln("Invalid input", input))
		}
	default:
		panic(fmt.Sprintln("Invalid state", e.state))
	}
}

func (e *Executor) Name() string {
	return e.processType + "." + e.processID
}
<|endoftext|>


<commit_before>// +build windows

package main

import (
	"fmt"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/windows/svc"

	"github.com/StackExchange/wmi"
	"github.com/martinlindhe/wmi_exporter/collector"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
	"gopkg.in/alecthomas/kingpin.v2"
)

// WmiCollector implements the prometheus.Collector interface.
type WmiCollector struct {
	maxScrapeDuration time.Duration
	collectors        map[string]collector.Collector
}

const (
	defaultCollectors            = "cpu,cs,logical_disk,net,os,service,system,textfile"
	defaultCollectorsPlaceholder = "[defaults]"
	serviceName                  = "wmi_exporter"
)

var (
	scrapeDurationDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_duration_seconds"),
		"wmi_exporter: Duration of a collection.",
		[]string{"collector"},
		nil,
	)
	scrapeSuccessDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_success"),
		"wmi_exporter: Whether the collector was successful.",
		[]string{"collector"},
		nil,
	)
	scrapeTimeoutDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_timeout"),
		"wmi_exporter: Whether the collector timed out.",
		[]string{"collector"},
		nil,
	)
	snapshotDuration = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "perflib_snapshot_duration_seconds"),
		"Duration of perflib snapshot capture",
		nil,
		nil,
	)

	// This can be removed when client_golang exposes this on Windows
	// (See https://github.com/prometheus/client_golang/issues/376)
	startTime     = float64(time.Now().Unix())
	startTimeDesc = prometheus.NewDesc(
		"process_start_time_seconds",
		"Start time of the process since unix epoch in seconds.",
		nil,
		nil,
	)
)

// Describe sends all the descriptors of the collectors included to
// the provided channel.
func (coll WmiCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- scrapeDurationDesc
	ch <- scrapeSuccessDesc
}

// Collect sends the collected metrics from each of the collectors to
// prometheus.
func (coll WmiCollector) Collect(ch chan<- prometheus.Metric) {
	t := time.Now()
	scrapeContext, err := collector.PrepareScrapeContext()
	ch <- prometheus.MustNewConstMetric(
		snapshotDuration,
		prometheus.GaugeValue,
		time.Since(t).Seconds(),
	)
	if err != nil {
		ch <- prometheus.NewInvalidMetric(scrapeSuccessDesc, fmt.Errorf("failed to prepare scrape: %v", err))
		return
	}

	remainingCollectors := make(map[string]bool)
	for name := range coll.collectors {
		remainingCollectors[name] = true
	}

	metricsBuffer := make(chan prometheus.Metric)
	allDone := make(chan struct{})
	stopped := false
	go func() {
		for {
			select {
			case m, ok := <-metricsBuffer:
				if ok && !stopped {
					ch <- m
				}
			case <-allDone:
				return
			}
		}
	}()

	wg := sync.WaitGroup{}
	wg.Add(len(coll.collectors))
	go func() {
		wg.Wait()
		close(allDone)
		close(metricsBuffer)
	}()

	for name, c := range coll.collectors {
		go func(name string, c collector.Collector) {
			defer wg.Done()
			execute(name, c, scrapeContext, metricsBuffer)
			delete(remainingCollectors, name)
		}(name, c)
	}

	ch <- prometheus.MustNewConstMetric(
		startTimeDesc,
		prometheus.CounterValue,
		startTime,
	)

	select {
	case <-allDone:
		stopped = true
		return
	case <-time.After(coll.maxScrapeDuration):
		stopped = true
		remainingCollectorNames := make([]string, 0, len(remainingCollectors))
		for rc := range remainingCollectors {
			remainingCollectorNames = append(remainingCollectorNames, rc)
		}
		log.Warn("Collection timed out, still waiting for ", remainingCollectorNames)
		for name := range remainingCollectors {
			ch <- prometheus.MustNewConstMetric(
				scrapeSuccessDesc,
				prometheus.GaugeValue,
				0.0,
				name,
			)
			ch <- prometheus.MustNewConstMetric(
				scrapeTimeoutDesc,
				prometheus.GaugeValue,
				1.0,
				name,
			)
		}
	}
}

func filterAvailableCollectors(collectors string) string {
	var availableCollectors []string
	for _, c := range strings.Split(collectors, ",") {
		_, ok := collector.Factories[c]
		if ok {
			availableCollectors = append(availableCollectors, c)
		}
	}
	return strings.Join(availableCollectors, ",")
}

func execute(name string, c collector.Collector, ctx *collector.ScrapeContext, ch chan<- prometheus.Metric) {
	begin := time.Now()
	err := c.Collect(ctx, ch)
	duration := time.Since(begin)
	var success float64

	if err != nil {
		log.Errorf("collector %s failed after %fs: %s", name, duration.Seconds(), err)
		success = 0
	} else {
		log.Debugf("collector %s succeeded after %fs.", name, duration.Seconds())
		success = 1
	}
	ch <- prometheus.MustNewConstMetric(
		scrapeDurationDesc,
		prometheus.GaugeValue,
		duration.Seconds(),
		name,
	)
	ch <- prometheus.MustNewConstMetric(
		scrapeSuccessDesc,
		prometheus.GaugeValue,
		success,
		name,
	)
	ch <- prometheus.MustNewConstMetric(
		scrapeTimeoutDesc,
		prometheus.GaugeValue,
		0.0,
		name,
	)
}

func expandEnabledCollectors(enabled string) []string {
	expanded := strings.Replace(enabled, defaultCollectorsPlaceholder, defaultCollectors, -1)
	separated := strings.Split(expanded, ",")
	unique := map[string]bool{}
	for _, s := range separated {
		if s != "" {
			unique[s] = true
		}
	}
	result := make([]string, 0, len(unique))
	for s := range unique {
		result = append(result, s)
	}
	return result
}

func loadCollectors(list string) (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	enabled := expandEnabledCollectors(list)

	for _, name := range enabled {
		fn, ok := collector.Factories[name]
		if !ok {
			return nil, fmt.Errorf("collector '%s' not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}
	return collectors, nil
}

func initWbem() {
	// This initialization prevents a memory leak on WMF 5+. See
	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
	// for details.
	log.Debugf("Initializing SWbemServices")
	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
	if err != nil {
		log.Fatal(err)
	}
	wmi.DefaultClient.AllowMissingFields = true
	wmi.DefaultClient.SWbemServicesClient = s
}

func main() {
	var (
		listenAddress = kingpin.Flag(
			"telemetry.addr",
			"host:port for WMI exporter.",
		).Default(":9182").String()
		metricsPath = kingpin.Flag(
			"telemetry.path",
			"URL path for surfacing collected metrics.",
		).Default("/metrics").String()
		enabledCollectors = kingpin.Flag(
			"collectors.enabled",
			"Comma-separated list of collectors to use. Use '[defaults]' as a placeholder for all the collectors enabled by default.").
			Default(filterAvailableCollectors(defaultCollectors)).String()
		printCollectors = kingpin.Flag(
			"collectors.print",
			"If true, print available collectors and exit.",
		).Bool()
		timeoutMargin = kingpin.Flag(
			"scrape.timeout-margin",
			"Seconds to subtract from the timeout allowed by the client. Tune to allow for overhead or high loads.",
		).Default("0.5").Float64()
	)

	log.AddFlags(kingpin.CommandLine)
	kingpin.Version(version.Print("wmi_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	if *printCollectors {
		collectorNames := make(sort.StringSlice, 0, len(collector.Factories))
		for n := range collector.Factories {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}

	initWbem()

	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Fatal(err)
	}

	stopCh := make(chan bool)
	if !isInteractive {
		go func() {
			err = svc.Run(serviceName, &wmiExporterService{stopCh: stopCh})
			if err != nil {
				log.Errorf("Failed to start service: %v", err)
			}
		}()
	}

	collectors, err := loadCollectors(*enabledCollectors)
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors: %v", strings.Join(keys(collectors), ", "))

	h := &metricsHandler{
		timeoutMargin: *timeoutMargin,
		collectorFactory: func(timeout time.Duration) *WmiCollector {
			return &WmiCollector{
				collectors:        collectors,
				maxScrapeDuration: timeout,
			}
		},
	}

	http.Handle(*metricsPath, h)
	http.HandleFunc("/health", healthCheck)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, *metricsPath, http.StatusMovedPermanently)
	})

	log.Infoln("Starting WMI exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	go func() {
		log.Infoln("Starting server on", *listenAddress)
		log.Fatalf("cannot start WMI exporter: %s", http.ListenAndServe(*listenAddress, nil))
	}()

	for {
		if <-stopCh {
			log.Info("Shutting down WMI exporter")
			break
		}
	}
}

func healthCheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_, err := fmt.Fprintln(w, `{"status":"ok"}`)
	if err != nil {
		log.Debugf("Failed to write to stream: %v", err)
	}
}

func keys(m map[string]collector.Collector) []string {
	ret := make([]string, 0, len(m))
	for key := range m {
		ret = append(ret, key)
	}
	return ret
}

type wmiExporterService struct {
	stopCh chan<- bool
}

func (s *wmiExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
loop:
	for {
		select {
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s.stopCh <- true
				break loop
			default:
				log.Error(fmt.Sprintf("unexpected control request #%d", c))
			}
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	return
}

type metricsHandler struct {
	timeoutMargin    float64
	collectorFactory func(timeout time.Duration) *WmiCollector
}

func (mh *metricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	const defaultTimeout = 10.0

	var timeoutSeconds float64
	if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
		var err error
		timeoutSeconds, err = strconv.ParseFloat(v, 64)
		if err != nil {
			log.Warnf("Couldn't parse X-Prometheus-Scrape-Timeout-Seconds: %q. Defaulting timeout to %f", v, defaultTimeout)
		}
	}
	if timeoutSeconds == 0 {
		timeoutSeconds = defaultTimeout
	}
	timeoutSeconds = timeoutSeconds - mh.timeoutMargin

	reg := prometheus.NewRegistry()
	reg.MustRegister(mh.collectorFactory(time.Duration(timeoutSeconds * float64(time.Second))))
	reg.MustRegister(
		prometheus.NewProcessCollector(os.Getpid(), ""),
		prometheus.NewGoCollector(),
		version.NewCollector("wmi_exporter"),
	)

	h := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
	h.ServeHTTP(w, r)
}
<commit_msg>Refactor timeout management<commit_after>// +build windows

package main

import (
	"fmt"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/windows/svc"

	"github.com/StackExchange/wmi"
	"github.com/martinlindhe/wmi_exporter/collector"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
	"gopkg.in/alecthomas/kingpin.v2"
)

// WmiCollector implements the prometheus.Collector interface.
type WmiCollector struct {
	maxScrapeDuration time.Duration
	collectors        map[string]collector.Collector
}

const (
	defaultCollectors            = "cpu,cs,logical_disk,net,os,service,system,textfile"
	defaultCollectorsPlaceholder = "[defaults]"
	serviceName                  = "wmi_exporter"
)

var (
	scrapeDurationDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_duration_seconds"),
		"wmi_exporter: Duration of a collection.",
		[]string{"collector"},
		nil,
	)
	scrapeSuccessDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_success"),
		"wmi_exporter: Whether the collector was successful.",
		[]string{"collector"},
		nil,
	)
	scrapeTimeoutDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_timeout"),
		"wmi_exporter: Whether the collector timed out.",
		[]string{"collector"},
		nil,
	)
	snapshotDuration = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "perflib_snapshot_duration_seconds"),
		"Duration of perflib snapshot capture",
		nil,
		nil,
	)

	// This can be removed when client_golang exposes this on Windows
	// (See https://github.com/prometheus/client_golang/issues/376)
	startTime     = float64(time.Now().Unix())
	startTimeDesc = prometheus.NewDesc(
		"process_start_time_seconds",
		"Start time of the process since unix epoch in seconds.",
		nil,
		nil,
	)
)

// Describe sends all the descriptors of the collectors included to
// the provided channel.
func (coll WmiCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- scrapeDurationDesc
	ch <- scrapeSuccessDesc
}

type collectorOutcome int

const (
	pending collectorOutcome = iota
	success
	failed
)

// Collect sends the collected metrics from each of the collectors to
// prometheus.
func (coll WmiCollector) Collect(ch chan<- prometheus.Metric) {
	ch <- prometheus.MustNewConstMetric(
		startTimeDesc,
		prometheus.CounterValue,
		startTime,
	)

	t := time.Now()
	scrapeContext, err := collector.PrepareScrapeContext()
	ch <- prometheus.MustNewConstMetric(
		snapshotDuration,
		prometheus.GaugeValue,
		time.Since(t).Seconds(),
	)
	if err != nil {
		ch <- prometheus.NewInvalidMetric(scrapeSuccessDesc, fmt.Errorf("failed to prepare scrape: %v", err))
		return
	}

	wg := sync.WaitGroup{}
	wg.Add(len(coll.collectors))
	collectorOutcomes := make(map[string]collectorOutcome)
	for name := range coll.collectors {
		collectorOutcomes[name] = pending
	}

	metricsBuffer := make(chan prometheus.Metric)
	l := sync.Mutex{}
	finished := false
	go func() {
		for m := range metricsBuffer {
			l.Lock()
			if !finished {
				ch <- m
			}
			l.Unlock()
		}
	}()

	for name, c := range coll.collectors {
		go func(name string, c collector.Collector) {
			defer wg.Done()
			outcome := execute(name, c, scrapeContext, metricsBuffer)
			l.Lock()
			if !finished {
				collectorOutcomes[name] = outcome
			}
			l.Unlock()
		}(name, c)
	}

	allDone := make(chan struct{})
	go func() {
		wg.Wait()
		close(allDone)
	}()

	// Wait until either all collectors finish, or timeout expires
	select {
	case <-allDone:
	case <-time.After(coll.maxScrapeDuration):
	}

	l.Lock()
	finished = true

	remainingCollectorNames := make([]string, 0)
	for name, outcome := range collectorOutcomes {
		var successValue, timeoutValue float64
		if outcome == pending {
			timeoutValue = 1.0
			remainingCollectorNames = append(remainingCollectorNames, name)
		}
		if outcome == success {
			successValue = 1.0
		}

		ch <- prometheus.MustNewConstMetric(
			scrapeSuccessDesc,
			prometheus.GaugeValue,
			successValue,
			name,
		)
		ch <- prometheus.MustNewConstMetric(
			scrapeTimeoutDesc,
			prometheus.GaugeValue,
			timeoutValue,
			name,
		)
	}

	if len(remainingCollectorNames) > 0 {
		log.Warn("Collection timed out, still waiting for ", remainingCollectorNames)
	}

	l.Unlock()
}

func filterAvailableCollectors(collectors string) string {
	var availableCollectors []string
	for _, c := range strings.Split(collectors, ",") {
		_, ok := collector.Factories[c]
		if ok {
			availableCollectors = append(availableCollectors, c)
		}
	}
	return strings.Join(availableCollectors, ",")
}

func execute(name string, c collector.Collector, ctx *collector.ScrapeContext, ch chan<- prometheus.Metric) collectorOutcome {
	t := time.Now()
	err := c.Collect(ctx, ch)
	duration := time.Since(t).Seconds()
	ch <- prometheus.MustNewConstMetric(
		scrapeDurationDesc,
		prometheus.GaugeValue,
		duration,
		name,
	)

	if err != nil {
		log.Errorf("collector %s failed after %fs: %s", name, duration, err)
		return failed
	}
	log.Debugf("collector %s succeeded after %fs.", name, duration)
	return success
}

func expandEnabledCollectors(enabled string) []string {
	expanded := strings.Replace(enabled, defaultCollectorsPlaceholder, defaultCollectors, -1)
	separated := strings.Split(expanded, ",")
	unique := map[string]bool{}
	for _, s := range separated {
		if s != "" {
			unique[s] = true
		}
	}
	result := make([]string, 0, len(unique))
	for s := range unique {
		result = append(result, s)
	}
	return result
}

func loadCollectors(list string) (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	enabled := expandEnabledCollectors(list)

	for _, name := range enabled {
		fn, ok := collector.Factories[name]
		if !ok {
			return nil, fmt.Errorf("collector '%s' not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}
	return collectors, nil
}

func initWbem() {
	// This initialization prevents a memory leak on WMF 5+. See
	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
	// for details.
	log.Debugf("Initializing SWbemServices")
	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
	if err != nil {
		log.Fatal(err)
	}
	wmi.DefaultClient.AllowMissingFields = true
	wmi.DefaultClient.SWbemServicesClient = s
}

func main() {
	var (
		listenAddress = kingpin.Flag(
			"telemetry.addr",
			"host:port for WMI exporter.",
		).Default(":9182").String()
		metricsPath = kingpin.Flag(
			"telemetry.path",
			"URL path for surfacing collected metrics.",
		).Default("/metrics").String()
		enabledCollectors = kingpin.Flag(
			"collectors.enabled",
			"Comma-separated list of collectors to use. Use '[defaults]' as a placeholder for all the collectors enabled by default.").
			Default(filterAvailableCollectors(defaultCollectors)).String()
		printCollectors = kingpin.Flag(
			"collectors.print",
			"If true, print available collectors and exit.",
		).Bool()
		timeoutMargin = kingpin.Flag(
			"scrape.timeout-margin",
			"Seconds to subtract from the timeout allowed by the client. Tune to allow for overhead or high loads.",
		).Default("0.5").Float64()
	)

	log.AddFlags(kingpin.CommandLine)
	kingpin.Version(version.Print("wmi_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	if *printCollectors {
		collectorNames := make(sort.StringSlice, 0, len(collector.Factories))
		for n := range collector.Factories {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}

	initWbem()

	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Fatal(err)
	}

	stopCh := make(chan bool)
	if !isInteractive {
		go func() {
			err = svc.Run(serviceName, &wmiExporterService{stopCh: stopCh})
			if err != nil {
				log.Errorf("Failed to start service: %v", err)
			}
		}()
	}

	collectors, err := loadCollectors(*enabledCollectors)
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors: %v", strings.Join(keys(collectors), ", "))

	h := &metricsHandler{
		timeoutMargin: *timeoutMargin,
		collectorFactory: func(timeout time.Duration) *WmiCollector {
			return &WmiCollector{
				collectors:        collectors,
				maxScrapeDuration: timeout,
			}
		},
	}

	http.Handle(*metricsPath, h)
	http.HandleFunc("/health", healthCheck)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, *metricsPath, http.StatusMovedPermanently)
	})

	log.Infoln("Starting WMI exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	go func() {
		log.Infoln("Starting server on", *listenAddress)
		log.Fatalf("cannot start WMI exporter: %s", http.ListenAndServe(*listenAddress, nil))
	}()

	for {
		if <-stopCh {
			log.Info("Shutting down WMI exporter")
			break
		}
	}
}

func healthCheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_, err := fmt.Fprintln(w, `{"status":"ok"}`)
	if err != nil {
		log.Debugf("Failed to write to stream: %v", err)
	}
}

func keys(m map[string]collector.Collector) []string {
	ret := make([]string, 0, len(m))
	for key := range m {
		ret = append(ret, key)
	}
	return ret
}

type wmiExporterService struct {
	stopCh chan<- bool
}

func (s *wmiExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
loop:
	for {
		select {
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s.stopCh <- true
				break loop
			default:
				log.Error(fmt.Sprintf("unexpected control request #%d", c))
			}
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	return
}

type metricsHandler struct {
	timeoutMargin    float64
	collectorFactory func(timeout time.Duration) *WmiCollector
}

func (mh *metricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	const defaultTimeout = 10.0

	var timeoutSeconds float64
	if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
		var err error
		timeoutSeconds, err = strconv.ParseFloat(v, 64)
		if err != nil {
			log.Warnf("Couldn't parse X-Prometheus-Scrape-Timeout-Seconds: %q. Defaulting timeout to %f", v, defaultTimeout)
		}
	}
	if timeoutSeconds == 0 {
		timeoutSeconds = defaultTimeout
	}
	timeoutSeconds = timeoutSeconds - mh.timeoutMargin

	reg := prometheus.NewRegistry()
	reg.MustRegister(mh.collectorFactory(time.Duration(timeoutSeconds * float64(time.Second))))
	reg.MustRegister(
		prometheus.NewProcessCollector(os.Getpid(), ""),
		prometheus.NewGoCollector(),
		version.NewCollector("wmi_exporter"),
	)

	h := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
	h.ServeHTTP(w, r)
}
<|endoftext|>