<commit_before>// Package oauth2 implements an OAuth 2.0 authorization server with support
// for token revokation.
//
// For details about the specs implemented please refer to
// * http://tools.ietf.org/html/rfc6749
// * http://tools.ietf.org/html/rfc6750
// * https://tools.ietf.org/html/rfc7009
package oauth2

import (
	"html/template"
	"log"
	"net/http"
	"strings"
	"time"
)

// Client defines client information required by oauth2 to:
//   * Show an authorization form to a resource owner
//   * Validate that the provided request_uri parameter matches the one previously
//     registered for the client.
type Client struct {
	ID            string
	Name          string
	Desc          string
	ProfileImgURL string
	HomepageURL   string
	RedirectURL   string
}

// Scope defines a type for manipulating OAuth2 scopes.
type Scope struct {
	ID   string
	Desc string
}

// Defines a type for the two defined token types in OAuth2.
type TokenType string

const (
	AccessToken  TokenType = "access"
	RefreshToken TokenType = "refresh"
)

type Provider interface {
	// ClientInfo returns 3rd-party client information
	ClientInfo(clientID string) (info Client, err error)

	// GenAuthzCode issues and stores an authorization grant code, in a persistent storage.
	GenAuthzCode(clientID, scopes []Scope) (code string, err error)

	// RevokeAuthzCode expires the grant code as well as all access and refresh tokens generated with it.
	RevokeAuthzCode(code string) error

	// ScopesInfo parses the list of scopes requested by the client and
	// returns its descriptions for the resource owner to fully understand
	// what he is authorizing the client to access to. An error is returned
	// if the scopes list does not comply with http://tools.ietf.org/html/rfc6749#section-3.3
	//
	// Unrecognized or non-existent scopes are ignored.
	ScopesInfo(scopes string) ([]Scope, error)

	// GenToken generates and stores token.
	GenToken(tokenType TokenType, scopes []Scope) (token string, err error)

	// RevokeToken expires a specific token.
	RevokeToken(token string) error

	// RefreshToken refreshes an access token.
	RefreshToken(refreshToken, scopes []Scope) (accessToken string, err error)

	// AuthzForm returns the HTML authorization form.
	AuthzForm() string

	// LoginURL returns the login URL for the resource owner to authenticate if there is
	// not a valid session. The authentication system should send back the user
	// to the referer URL in order to complete the OAuth2 authorization process.
	LoginURL(refererURL string) (url string)

	// CheckSession checks whether or not the resource owner has a valid session
	// with the system. If not, it redirects the user to the login URL.
	CheckSession() (invalid bool)
}

// http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type option func(*config)

type config struct {
	authzEndpoint  string
	tokenEndpoint  string
	revokeEndpoint string
	stsMaxAge      time.Duration
	authzForm      *template.Template
	provider       Provider
}

// TokenEndpoint allows setting token endpoint. Defaults to "/oauth2/tokens".
//
// The token endpoint is used by the client to obtain an access token by
// presenting its authorization grant or refresh token.  The token
// endpoint is used with every authorization grant except for the
// implicit grant type (since an access token is issued directly).
//
// Since requests to the token endpoint result in the transmission of
// clear-text credentials (in the HTTP request and response), the
// authorization server MUST require the use of TLS as described in
// Section 1.6 when sending requests to the token endpoint.
//
// -- http://tools.ietf.org/html/rfc6749#section-3.2
func SetTokenEndpoint(endpoint string) option {
	return func(c *config) {
		c.tokenEndpoint = endpoint
	}
}

// AuthzEndpoint allows setting authorization endpoint. Defaults to "/oauth2/authzs"
//
// The authorization endpoint is used to interact with the resource owner and
// obtain an authorization grant.
//
// Since requests to the authorization endpoint result in user authentication
// and the transmission of clear-text credentials (in the HTTP response), the
// authorization server MUST require the use of TLS as described in Section 1.6
// when sending requests to the authorization endpoint.
//
// -- http://tools.ietf.org/html/rfc6749#section-3.1.1
func SetAuthzEndpoint(endpoint string) option {
	return func(c *config) {
		c.authzEndpoint = endpoint
	}
}

// SetRevokeEndpoint allows setting a custom token revoke URI. Defaults to "/oauth2/revoke"
func SetRevokeEndpoint(endpoint string) option {
	return func(c *config) {
		c.revokeEndpoint = endpoint
	}
}

// SetSTSMaxAge sets Strict Transport Security maximum age. Defaults to 1yr
func SetSTSMaxAge(maxAge time.Duration) option {
	return func(c *config) {
		c.stsMaxAge = maxAge
	}
}

// SetAuthzForm sets authorization form to show to the resource owner
func SetAuthzForm(form string) option {
	return func(c *config) {
		t := template.New("authzform")
		tpl, err := t.Parse(form)
		if err != nil {
			log.Fatalln("Error parsing authorization form: %v", err)
		}
		c.authzForm = tpl
	}
}

// Handler handles OAuth2 requests.
func Handler(next http.Handler, opts ...option) http.Handler {
	// Default configuration options.
	cfg := &config{
		tokenEndpoint:  "/oauth2/tokens",
		authzEndpoint:  "/oauth2/authzs",
		revokeEndpoint: "/oauth2/revoke",
	}

	cfg.stsMaxAge = time.Duration(31536000) * time.Second // 1yr

	// Applies user's configuration.
	for _, opt := range opts {
		opt(cfg)
	}

	if cfg.authzForm == nil {
		log.Fatalln("Authorization form is required")
	}

	if cfg.provider == nil {
		log.Fatalln("An implementation of the oauth2.Provider interface is expected")
	}

	// Keeps a registry of path function handlers for OAuth2 requests.
	registry := map[string]map[string]func(http.ResponseWriter, *http.Request, *config, http.Handler){
		cfg.authzEndpoint: AuthzHandlers,
		cfg.tokenEndpoint: TokenHandlers,
		// TODO(c4milo): URL handlers for revoking tokens and grants
	}

	// Locates and runs specific OAuth2 handler for request's method
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		for p, handlers := range registry {
			if strings.HasPrefix(req.URL.Path, p) {
				if handlerFn, ok := handlers[req.Method]; ok {
					handlerFn(w, req, cfg, next)
					return
				}
				w.WriteHeader(http.StatusMethodNotAllowed)
				w.Write([]byte("Method Not Allowed"))
				return
			}
		}
	})
}
<commit_msg>Documents Provider interface<commit_after>// Package oauth2 implements an OAuth 2.0 authorization server with support
// for token revokation.
//
// For details about the specs implemented please refer to
// * http://tools.ietf.org/html/rfc6749
// * http://tools.ietf.org/html/rfc6750
// * https://tools.ietf.org/html/rfc7009
package oauth2

import (
	"html/template"
	"log"
	"net/http"
	"strings"
	"time"
)

// Client defines client information required by oauth2 to:
//   * Show an authorization form to a resource owner
//   * Validate that the provided request_uri parameter matches the one previously
//     registered for the client.
type Client struct {
	ID            string
	Name          string
	Desc          string
	ProfileImgURL string
	HomepageURL   string
	RedirectURL   string
}

// Scope defines a type for manipulating OAuth2 scopes.
type Scope struct {
	ID   string
	Desc string
}

// Defines a type for the two defined token types in OAuth2.
type TokenType string

const (
	AccessToken  TokenType = "access"
	RefreshToken TokenType = "refresh"
)

// Provider defines functions required by the oauth2 package to properly work.
// Users of this package are required to implement them.
type Provider interface {
	// ClientInfo returns 3rd-party client information
	ClientInfo(clientID string) (info Client, err error)

	// GenAuthzCode issues and stores an authorization grant code, in a persistent storage.
	GenAuthzCode(clientID, scopes []Scope) (code string, err error)

	// RevokeAuthzCode expires the grant code as well as all access and refresh tokens generated with it.
	RevokeAuthzCode(code string) error

	// ScopesInfo parses the list of scopes requested by the client and
	// returns its descriptions for the resource owner to fully understand
	// what he is authorizing the client to access to. An error is returned
	// if the scopes list does not comply with http://tools.ietf.org/html/rfc6749#section-3.3
	//
	// Unrecognized or non-existent scopes are ignored.
	ScopesInfo(scopes string) ([]Scope, error)

	// GenToken generates and stores token.
	GenToken(tokenType TokenType, scopes []Scope) (token string, err error)

	// RevokeToken expires a specific token.
	RevokeToken(token string) error

	// RefreshToken refreshes an access token.
	RefreshToken(refreshToken, scopes []Scope) (accessToken string, err error)

	// AuthzForm returns the HTML authorization form.
	AuthzForm() string

	// LoginURL returns the login URL for the resource owner to authenticate if there is
	// not a valid session. The authentication system should send back the user
	// to the referer URL in order to complete the OAuth2 authorization process.
	LoginURL(refererURL string) (url string)

	// CheckSession checks whether or not the resource owner has a valid session
	// with the system. If not, it redirects the user to the login URL.
	CheckSession() (invalid bool)
}

// http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type option func(*config)

type config struct {
	authzEndpoint  string
	tokenEndpoint  string
	revokeEndpoint string
	stsMaxAge      time.Duration
	authzForm      *template.Template
	provider       Provider
}

// TokenEndpoint allows setting token endpoint. Defaults to "/oauth2/tokens".
//
// The token endpoint is used by the client to obtain an access token by
// presenting its authorization grant or refresh token.  The token
// endpoint is used with every authorization grant except for the
// implicit grant type (since an access token is issued directly).
//
// Since requests to the token endpoint result in the transmission of
// clear-text credentials (in the HTTP request and response), the
// authorization server MUST require the use of TLS as described in
// Section 1.6 when sending requests to the token endpoint.
//
// -- http://tools.ietf.org/html/rfc6749#section-3.2
func SetTokenEndpoint(endpoint string) option {
	return func(c *config) {
		c.tokenEndpoint = endpoint
	}
}

// AuthzEndpoint allows setting authorization endpoint. Defaults to "/oauth2/authzs"
//
// The authorization endpoint is used to interact with the resource owner and
// obtain an authorization grant.
//
// Since requests to the authorization endpoint result in user authentication
// and the transmission of clear-text credentials (in the HTTP response), the
// authorization server MUST require the use of TLS as described in Section 1.6
// when sending requests to the authorization endpoint.
//
// -- http://tools.ietf.org/html/rfc6749#section-3.1.1
func SetAuthzEndpoint(endpoint string) option {
	return func(c *config) {
		c.authzEndpoint = endpoint
	}
}

// SetRevokeEndpoint allows setting a custom token revoke URI. Defaults to "/oauth2/revoke"
func SetRevokeEndpoint(endpoint string) option {
	return func(c *config) {
		c.revokeEndpoint = endpoint
	}
}

// SetSTSMaxAge sets Strict Transport Security maximum age. Defaults to 1yr
func SetSTSMaxAge(maxAge time.Duration) option {
	return func(c *config) {
		c.stsMaxAge = maxAge
	}
}

// SetAuthzForm sets authorization form to show to the resource owner
func SetAuthzForm(form string) option {
	return func(c *config) {
		t := template.New("authzform")
		tpl, err := t.Parse(form)
		if err != nil {
			log.Fatalln("Error parsing authorization form: %v", err)
		}
		c.authzForm = tpl
	}
}

// Handler handles OAuth2 requests.
func Handler(next http.Handler, opts ...option) http.Handler {
	// Default configuration options.
	cfg := &config{
		tokenEndpoint:  "/oauth2/tokens",
		authzEndpoint:  "/oauth2/authzs",
		revokeEndpoint: "/oauth2/revoke",
	}

	cfg.stsMaxAge = time.Duration(31536000) * time.Second // 1yr

	// Applies user's configuration.
	for _, opt := range opts {
		opt(cfg)
	}

	if cfg.authzForm == nil {
		log.Fatalln("Authorization form is required")
	}

	if cfg.provider == nil {
		log.Fatalln("An implementation of the oauth2.Provider interface is expected")
	}

	// Keeps a registry of path function handlers for OAuth2 requests.
	registry := map[string]map[string]func(http.ResponseWriter, *http.Request, *config, http.Handler){
		cfg.authzEndpoint: AuthzHandlers,
		cfg.tokenEndpoint: TokenHandlers,
		// TODO(c4milo): URL handlers for revoking tokens and grants
	}

	// Locates and runs specific OAuth2 handler for request's method
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		for p, handlers := range registry {
			if strings.HasPrefix(req.URL.Path, p) {
				if handlerFn, ok := handlers[req.Method]; ok {
					handlerFn(w, req, cfg, next)
					return
				}
				w.WriteHeader(http.StatusMethodNotAllowed)
				w.Write([]byte("Method Not Allowed"))
				return
			}
		}
	})
}
<|endoftext|>


<commit_before>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * promise/functions.go                                   *
 *                                                        *
 * some functions of promise for Go.                      *
 *                                                        *
 * LastModified: Aug 13, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package promise

import (
	"errors"
	"sync/atomic"

	"github.com/hprose/hprose-golang/promise"
)

func allHandler(promise Promise, count *int64, result []interface{}, value interface{}, index int) {
	ToPromise(value).Then(func(value interface{}) {
		result[index] = value
		if atomic.AddInt64(count, -1) == 0 {
			promise.Resolve(result)
		}
	}, promise.Reject)
}

// All function returns a promise that resolves when all of the promises in the
// iterable argument have resolved, or rejects with the reason of the first
// passed promise that rejects.
func All(iterable ...interface{}) Promise {
	count := int64(len(iterable))
	result := make([]interface{}, count)
	if count == 0 {
		return Resolve(result)
	}
	promise := New()
	for index, value := range iterable {
		allHandler(promise, &count, result, value, index)
	}
	return promise
}

// Race function returns a promise that resolves or rejects as soon as one of
// the promises in the iterable resolves or rejects, with the value or reason
// from that promise.
func Race(iterable ...interface{}) Promise {
	promise := New()
	for _, value := range iterable {
		ToPromise(value).Fill(promise)
	}
	return promise
}

func getAnyOnRejected(count *int64) {
	if atomic.AddInt64(count, -1) == 0 {
		promise.Reject(errors.New("any(): all promises failed"))
	}
}

// Any function is a competitive race that allows one winner.
//
// The returned promise will:
//
//     fulfill as soon as any one of the input promises fulfills, with the
//     value of the fulfilled input promise, or
//
//     reject:
//
//         with a IllegalArgumentError if the input array is empty--i.e.
//		   it is impossible to have one winner.
//
//         with an array of all the rejection reasons, if the input array is
//         non-empty, and all input promises reject.
func Any(iterable ...interface{}) Promise {
	count := int64(len(iterable))
	if count == 0 {
		return Reject(IllegalArgumentError("any(): array must not be empty"))
	}
	promise := New()
	for _, value := range iterable {
		ToPromise(value).Then(promise.Resolve, getAnyOnRejected(&count))
	}
	return promise
}
<commit_msg>Fixed Any function<commit_after>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * promise/functions.go                                   *
 *                                                        *
 * some functions of promise for Go.                      *
 *                                                        *
 * LastModified: Aug 13, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package promise

import (
	"errors"
	"sync/atomic"

	"github.com/hprose/hprose-golang/promise"
)

func allHandler(promise Promise, count *int64, result []interface{}, value interface{}, index int) {
	ToPromise(value).Then(func(value interface{}) {
		result[index] = value
		if atomic.AddInt64(count, -1) == 0 {
			promise.Resolve(result)
		}
	}, promise.Reject)
}

// All function returns a promise that resolves when all of the promises in the
// iterable argument have resolved, or rejects with the reason of the first
// passed promise that rejects.
func All(iterable ...interface{}) Promise {
	count := int64(len(iterable))
	result := make([]interface{}, count)
	if count == 0 {
		return Resolve(result)
	}
	promise := New()
	for index, value := range iterable {
		allHandler(promise, &count, result, value, index)
	}
	return promise
}

// Race function returns a promise that resolves or rejects as soon as one of
// the promises in the iterable resolves or rejects, with the value or reason
// from that promise.
func Race(iterable ...interface{}) Promise {
	promise := New()
	for _, value := range iterable {
		ToPromise(value).Fill(promise)
	}
	return promise
}

func getAnyOnRejected(count *int64) {
	return func() {
		if atomic.AddInt64(count, -1) == 0 {
			promise.Reject(errors.New("any(): all promises failed"))
		}
	}
}

// Any function is a competitive race that allows one winner.
//
// The returned promise will:
//
//     fulfill as soon as any one of the input promises fulfills, with the
//     value of the fulfilled input promise, or
//
//     reject:
//
//         with a IllegalArgumentError if the input array is empty--i.e.
//		   it is impossible to have one winner.
//
//         with an array of all the rejection reasons, if the input array is
//         non-empty, and all input promises reject.
func Any(iterable ...interface{}) Promise {
	count := int64(len(iterable))
	if count == 0 {
		return Reject(IllegalArgumentError("any(): array must not be empty"))
	}
	promise := New()
	for _, value := range iterable {
		ToPromise(value).Then(promise.Resolve, getAnyOnRejected(&count))
	}
	return promise
}
<|endoftext|>


<commit_before>// +build !windows

package system

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"runtime"
	"strconv"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

type Processes struct {
	execPS       func() ([]byte, error)
	readProcFile func(statFile string) ([]byte, error)

	forcePS   bool
	forceProc bool
}

func (p *Processes) Description() string {
	return "Get the number of processes and group them by status"
}

func (p *Processes) SampleConfig() string { return "" }

func (p *Processes) Gather(acc telegraf.Accumulator) error {
	// Get an empty map of metric fields
	fields := getEmptyFields()

	// Decide if we will use 'ps' to get stats (use procfs otherwise)
	usePS := true
	if runtime.GOOS == "linux" {
		usePS = false
	}
	if p.forcePS {
		usePS = true
	} else if p.forceProc {
		usePS = false
	}

	// Gather stats from 'ps' or procfs
	if usePS {
		if err := p.gatherFromPS(fields); err != nil {
			return err
		}
	} else {
		if err := p.gatherFromProc(fields); err != nil {
			return err
		}
	}

	acc.AddFields("processes", fields, nil)
	return nil
}

// Gets empty fields of metrics based on the OS
func getEmptyFields() map[string]interface{} {
	fields := map[string]interface{}{
		"blocked":  int64(0),
		"zombies":  int64(0),
		"stopped":  int64(0),
		"running":  int64(0),
		"sleeping": int64(0),
		"total":    int64(0),
	}
	switch runtime.GOOS {
	case "freebsd":
		fields["idle"] = int64(0)
		fields["wait"] = int64(0)
	case "darwin":
		fields["idle"] = int64(0)
	case "openbsd":
		fields["idle"] = int64(0)
	case "linux":
		fields["paging"] = int64(0)
		fields["total_threads"] = int64(0)
	}
	return fields
}

// exec `ps` to get all process states
func (p *Processes) gatherFromPS(fields map[string]interface{}) error {
	out, err := p.execPS()
	if err != nil {
		return err
	}

	for i, status := range bytes.Fields(out) {
		if i == 0 && string(status) == "STAT" {
			// This is a header, skip it
			continue
		}
		switch status[0] {
		case 'W':
			fields["wait"] = fields["wait"].(int64) + int64(1)
		case 'U', 'D', 'L':
			// Also known as uninterruptible sleep or disk sleep
			fields["blocked"] = fields["blocked"].(int64) + int64(1)
		case 'Z':
			fields["zombies"] = fields["zombies"].(int64) + int64(1)
		case 'T':
			fields["stopped"] = fields["stopped"].(int64) + int64(1)
		case 'R':
			fields["running"] = fields["running"].(int64) + int64(1)
		case 'S':
			fields["sleeping"] = fields["sleeping"].(int64) + int64(1)
		case 'I':
			fields["idle"] = fields["idle"].(int64) + int64(1)
		default:
			log.Printf("processes: Unknown state [ %s ] from ps",
				string(status[0]))
		}
		fields["total"] = fields["total"].(int64) + int64(1)
	}
	return nil
}

// get process states from /proc/(pid)/stat files
func (p *Processes) gatherFromProc(fields map[string]interface{}) error {
	files, err := ioutil.ReadDir("/proc")
	if err != nil {
		return err
	}

	for _, file := range files {
		if !file.IsDir() {
			continue
		}

		statFile := path.Join("/proc", file.Name(), "stat")
		data, err := p.readProcFile(statFile)
		if err != nil {
			return err
		}
		if data == nil {
			continue
		}

		// Parse out data after (<cmd name>)
		i := bytes.LastIndex(data, []byte(")"))
		if i == -1 {
			continue
		}
		data = data[i+2:]

		stats := bytes.Fields(data)
		if len(stats) < 3 {
			return fmt.Errorf("Something is terribly wrong with %s", statFile)
		}
		switch stats[0][0] {
		case 'R':
			fields["running"] = fields["running"].(int64) + int64(1)
		case 'S':
			fields["sleeping"] = fields["sleeping"].(int64) + int64(1)
		case 'D':
			fields["blocked"] = fields["blocked"].(int64) + int64(1)
		case 'Z':
			fields["zombies"] = fields["zombies"].(int64) + int64(1)
		case 'T', 't':
			fields["stopped"] = fields["stopped"].(int64) + int64(1)
		case 'W':
			fields["paging"] = fields["paging"].(int64) + int64(1)
		default:
			log.Printf("processes: Unknown state [ %s ] in file %s",
				string(stats[0][0]), statFile)
		}
		fields["total"] = fields["total"].(int64) + int64(1)

		threads, err := strconv.Atoi(string(stats[17]))
		if err != nil {
			log.Printf("processes: Error parsing thread count: %s", err)
			continue
		}
		fields["total_threads"] = fields["total_threads"].(int64) + int64(threads)
	}
	return nil
}

func readProcFile(statFile string) ([]byte, error) {
	if _, err := os.Stat(statFile); os.IsNotExist(err) {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	data, err := ioutil.ReadFile(statFile)
	if err != nil {
		return nil, err
	}

	return data, nil
}

func execPS() ([]byte, error) {
	bin, err := exec.LookPath("ps")
	if err != nil {
		return nil, err
	}

	out, err := exec.Command(bin, "axo", "state").Output()
	if err != nil {
		return nil, err
	}

	return out, err
}

func init() {
	inputs.Add("processes", func() telegraf.Input {
		return &Processes{
			execPS:       execPS,
			readProcFile: readProcFile,
		}
	})
}
<commit_msg>processes: add 'unknown' procs (?)<commit_after>// +build !windows

package system

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"runtime"
	"strconv"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

type Processes struct {
	execPS       func() ([]byte, error)
	readProcFile func(statFile string) ([]byte, error)

	forcePS   bool
	forceProc bool
}

func (p *Processes) Description() string {
	return "Get the number of processes and group them by status"
}

func (p *Processes) SampleConfig() string { return "" }

func (p *Processes) Gather(acc telegraf.Accumulator) error {
	// Get an empty map of metric fields
	fields := getEmptyFields()

	// Decide if we will use 'ps' to get stats (use procfs otherwise)
	usePS := true
	if runtime.GOOS == "linux" {
		usePS = false
	}
	if p.forcePS {
		usePS = true
	} else if p.forceProc {
		usePS = false
	}

	// Gather stats from 'ps' or procfs
	if usePS {
		if err := p.gatherFromPS(fields); err != nil {
			return err
		}
	} else {
		if err := p.gatherFromProc(fields); err != nil {
			return err
		}
	}

	acc.AddFields("processes", fields, nil)
	return nil
}

// Gets empty fields of metrics based on the OS
func getEmptyFields() map[string]interface{} {
	fields := map[string]interface{}{
		"blocked":  int64(0),
		"zombies":  int64(0),
		"stopped":  int64(0),
		"running":  int64(0),
		"sleeping": int64(0),
		"total":    int64(0),
		"unknown":  int64(0),
	}
	switch runtime.GOOS {
	case "freebsd":
		fields["idle"] = int64(0)
		fields["wait"] = int64(0)
	case "darwin":
		fields["idle"] = int64(0)
	case "openbsd":
		fields["idle"] = int64(0)
	case "linux":
		fields["paging"] = int64(0)
		fields["total_threads"] = int64(0)
	}
	return fields
}

// exec `ps` to get all process states
func (p *Processes) gatherFromPS(fields map[string]interface{}) error {
	out, err := p.execPS()
	if err != nil {
		return err
	}

	for i, status := range bytes.Fields(out) {
		if i == 0 && string(status) == "STAT" {
			// This is a header, skip it
			continue
		}
		switch status[0] {
		case 'W':
			fields["wait"] = fields["wait"].(int64) + int64(1)
		case 'U', 'D', 'L':
			// Also known as uninterruptible sleep or disk sleep
			fields["blocked"] = fields["blocked"].(int64) + int64(1)
		case 'Z':
			fields["zombies"] = fields["zombies"].(int64) + int64(1)
		case 'T':
			fields["stopped"] = fields["stopped"].(int64) + int64(1)
		case 'R':
			fields["running"] = fields["running"].(int64) + int64(1)
		case 'S':
			fields["sleeping"] = fields["sleeping"].(int64) + int64(1)
		case 'I':
			fields["idle"] = fields["idle"].(int64) + int64(1)
		case '?':
			fields["unknown"] = fields["unknown"].(int64) + int64(1)
		default:
			log.Printf("processes: Unknown state [ %s ] from ps",
				string(status[0]))
		}
		fields["total"] = fields["total"].(int64) + int64(1)
	}
	return nil
}

// get process states from /proc/(pid)/stat files
func (p *Processes) gatherFromProc(fields map[string]interface{}) error {
	files, err := ioutil.ReadDir("/proc")
	if err != nil {
		return err
	}

	for _, file := range files {
		if !file.IsDir() {
			continue
		}

		statFile := path.Join("/proc", file.Name(), "stat")
		data, err := p.readProcFile(statFile)
		if err != nil {
			return err
		}
		if data == nil {
			continue
		}

		// Parse out data after (<cmd name>)
		i := bytes.LastIndex(data, []byte(")"))
		if i == -1 {
			continue
		}
		data = data[i+2:]

		stats := bytes.Fields(data)
		if len(stats) < 3 {
			return fmt.Errorf("Something is terribly wrong with %s", statFile)
		}
		switch stats[0][0] {
		case 'R':
			fields["running"] = fields["running"].(int64) + int64(1)
		case 'S':
			fields["sleeping"] = fields["sleeping"].(int64) + int64(1)
		case 'D':
			fields["blocked"] = fields["blocked"].(int64) + int64(1)
		case 'Z':
			fields["zombies"] = fields["zombies"].(int64) + int64(1)
		case 'T', 't':
			fields["stopped"] = fields["stopped"].(int64) + int64(1)
		case 'W':
			fields["paging"] = fields["paging"].(int64) + int64(1)
		default:
			log.Printf("processes: Unknown state [ %s ] in file %s",
				string(stats[0][0]), statFile)
		}
		fields["total"] = fields["total"].(int64) + int64(1)

		threads, err := strconv.Atoi(string(stats[17]))
		if err != nil {
			log.Printf("processes: Error parsing thread count: %s", err)
			continue
		}
		fields["total_threads"] = fields["total_threads"].(int64) + int64(threads)
	}
	return nil
}

func readProcFile(statFile string) ([]byte, error) {
	if _, err := os.Stat(statFile); os.IsNotExist(err) {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	data, err := ioutil.ReadFile(statFile)
	if err != nil {
		return nil, err
	}

	return data, nil
}

func execPS() ([]byte, error) {
	bin, err := exec.LookPath("ps")
	if err != nil {
		return nil, err
	}

	out, err := exec.Command(bin, "axo", "state").Output()
	if err != nil {
		return nil, err
	}

	return out, err
}

func init() {
	inputs.Add("processes", func() telegraf.Input {
		return &Processes{
			execPS:       execPS,
			readProcFile: readProcFile,
		}
	})
}
<|endoftext|>


<commit_before>/*
 * Copyright 2012 Nan Deng
 *
 * 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 server

import (
	"github.com/uniqush/uniqush-conn/proto"
	"net"
)

type Conn interface {
	proto.Conn
}

type serverConn struct {
	proto.Conn
	cmdio *proto.CommandIO
}

func (self *serverConn) ProcessCommand(cmd *proto.Command) error {
	return nil
}

func NewConn(cmdio *proto.CommandIO, service, username string, conn net.Conn) Conn {
	sc := new(serverConn)
	sc.cmdio = cmdio
	c := proto.NewConn(cmdio, service, username, conn, sc)
	sc.Conn = c
	return sc
}

<commit_msg>Overload the WriteMessage method<commit_after>/*
 * Copyright 2012 Nan Deng
 *
 * 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 server

import (
	"github.com/uniqush/uniqush-conn/proto"
	"net"
)

type Conn interface {
	proto.Conn
}

type serverConn struct {
	proto.Conn
	cmdio *proto.CommandIO
}

func (self *serverConn) WriteMessage(msg *proto.Message, compress, encrypt bool) error {
	return self.Conn.WriteMessage(msg, compress, encrypt)
}

func (self *serverConn) ProcessCommand(cmd *proto.Command) error {
	return nil
}

func NewConn(cmdio *proto.CommandIO, service, username string, conn net.Conn) Conn {
	sc := new(serverConn)
	sc.cmdio = cmdio
	c := proto.NewConn(cmdio, service, username, conn, sc)
	sc.Conn = c
	return sc
}

<|endoftext|>


<commit_before>package puppetdb

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// Client This represents a connection to your puppetdb instance
type Client struct {
	BaseURL    string
	Cert       string
	Key        string
	httpClient *http.Client
	verbose    bool
}

// EventCountJSON A json object holding the results of a query to the eventcount api
type EventCountJSON struct {
	SubjectType string            `json:"subject-type"`
	Subject     map[string]string `json:"subject"`
	Failure     int64             `json:"failures"`
	Successes   int64             `json:"successes"`
	Noops       int64             `json:"noops"`
	Skips       int64             `json:"skips"`
}

// EventJSON A json object holding the results of a query to the event api.
type EventJSON struct {
	CertName          string `json:"certname"`
	OldValue          string `json:"old-value"`
	Property          string `json:"message"`
	Timestamp         string `json:"timestamp"`
	ResourceType      string `json:"resource-typ"`
	ResourceTitle     string `json:"resource-title"`
	NewValue          string `json:"new-value"`
	Message           string `json:"message"`
	Report            string `json:"report"`
	Status            string `json:"status"`
	File              string `json:"file"`
	ContainmentPath   string `json:"containment-path"`
	ContainmentClass  string `json:"containing-class"`
	RunStartTime      string `json:"run-start-time"`
	RunEndTime        string `json:"run-end-time"`
	ReportReceiveTime string `json:"report-receive-time"`
}

// FactJSON A json object holding the results of a query to the facts api.
type FactJSON struct {
	CertName    string `json:"certname"`
	Environment string `json:"environment"`
	Name        string `json:"name"`
	Value       string `json:"value"`
}

// NodeJSON A json object holding the results of query to the node api.
type NodeJSON struct {
	Certname                     string `json:"certname"`
	Deactivated                  string `json:"deactivated"`
	CatalogTimestamp             string `json:"catalog_timestamp"`
	FactsTimestamp               string `json:"facts_timestamp"`
	CatalogEnvironment           string `json:"catalog_environment"`
	FactsEnvironment             string `json:"facts_environment"`
	LatestReportHash             string `json:"latest_report_hash"`
	CachedCatalogStatus          string `json:"cached_catalog_status"`
	ReportEnvironment            string `json:"report_environment"`
	ReportTimestamp              string `json:"report_timestamp"`
	LatestReportCorrectiveChange string `json:"latest_report_corrective_change"`
	LatestReportNoop             bool   `json:"latest_report_noop"`
	LatestReportNoopPending      bool   `json:"latest_report_noop_pending"`
	Expired                      string `json:"expired"`
	LatestReportJobID            string `json:"latest_report_job_id"`
	LatestReportStatus           string `json:"latest_report_status"`
}

// PuppetdbVersion a simple struct holding the puppetdb version.
type PuppetdbVersion struct {
	Version string `json:"version"`
}

// ReportJSON A json abject holding the data for a query from the report api.
type ReportJSON struct {
	CertName             string `json:"certname"`
	PuppetVersion        string `json:"puppet-version"`
	Value                string `json:"value"`
	Hash                 string `json:"hash"`
	ReportFormat         int64  `json:"report-format"`
	ConfigurationVersion string `json:"configuration-version"`
	TransactionUUID      string `json:"transaction-uuid"`
	StartTime            string `json:"start-time"`
	EndTime              string `json:"end-time"`
	ReceiveTime          string `json:"receive-time"`
}

//Resource contains information about a puppet resource.
type Resource struct {
	Paramaters map[string]interface{} `json:"parameters"`
	Line       int                    `json:"line,omitempty"`
	Exported   bool                   `json:"exported,omitempty"`
	Tags       []string               `json:"tags,omitempty"`
	Title      string                 `json:"title,omitempty"`
	Type       string                 `json:"type,omitempty"`
	Resource   string                 `json:"resource,omitempty"`
	Certname   string                 `json:"certname,omitempty"`
}

// ValueMetricJSON A simple structholding a float value.
type ValueMetricJSON struct {
	Value float64
}

// getURL return the address of the puppetdb instance.
func getURL(host string, port int, ssl bool) string {
	if ssl {
		return fmt.Sprintf("https://%s:%v", host, port)
	} else {
		return fmt.Sprintf("http://%s:%v", host, port)
	}
}

// NewClient returns a http connection for your puppetdb instance.
func NewClient(host string, port int, verbose bool) *Client {

	tr := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	client := &http.Client{Transport: tr}
	return &Client{getURL(host, port, false), "", "", client, verbose}
}

// NewClientSSL returns a https connection for your puppetdb instance.
func NewClientSSL(host string, port int, key string, cert string, verbose bool) *Client {
	flag.Parse()
	cert2, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		log.Fatal(err)
	}
	tlsConfig := &tls.Config{
		InsecureSkipVerify: true,
		Certificates:       []tls.Certificate{cert2},
	}
	tlsConfig.BuildNameToCertificate()
	transport := &http.Transport{TLSClientConfig: tlsConfig}
	client := &http.Client{Transport: transport}
	return &Client{getURL(host, port, true), cert, key, client, verbose}

}

// NewClientTimeout returns a http connection for your puppetdb instance with a timeout.
func NewClientTimeout(host string, port int, verbose bool, timeout int) *Client {

	tr := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	client := &http.Client{Transport: tr, Timeout: time.Duration(timeout) * time.Second}
	return &Client{getURL(host, port, false), "", "", client, verbose}
}

// NewClientTimeoutSSL returns a http connection for your puppetdb instance with a timeout and ssl configured.
func NewClientTimeoutSSL(host string, port int, key string, cert string, verbose bool, timeout int) *Client {
	flag.Parse()
	cert2, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		log.Fatal(err)
	}
	tlsConfig := &tls.Config{
		InsecureSkipVerify: true,
		Certificates:       []tls.Certificate{cert2},
	}
	tlsConfig.BuildNameToCertificate()
	transport := &http.Transport{TLSClientConfig: tlsConfig}
	client := &http.Client{Transport: transport, Timeout: time.Duration(timeout) * time.Second}
	return &Client{getURL(host, port, true), cert, key, client, verbose}

}

// Get gets the given url and retrusn the result. In form of the given interface.
func (c *Client) Get(v interface{}, path string, params map[string]string) error {
	pathAndParams := path
	//TODO: Improve this
	if params != nil && len(params) > 0 {
		if !strings.Contains(path, "?") {
			pathAndParams += "?"
		}
		for k, v := range params {
			pathAndParams += fmt.Sprintf("%s=%s&", k, url.QueryEscape(v))
		}
	}
	resp, err := c.httpGet(pathAndParams)
	if err != nil {
		log.Print(err)
		return err
	}
	defer resp.Body.Close()
	if err != nil {
		log.Print(err)
		return err
	}
	json.NewDecoder(resp.Body).Decode(&v)
	return err
}

// Nodes Polls the nodes api of your puppetdb and returns the results in form of the NodeJSON type.
func (c *Client) Nodes() ([]NodeJSON, error) {
	ret := []NodeJSON{}
	err := c.Get(&ret, "nodes", nil)
	return ret, err
}

func (c *Client) FactNames() ([]string, error) {
	ret := []string{}
	err := c.Get(&ret, "fact-names", nil)
	return ret, err
}

func (c *Client) NodeFacts(node string) ([]FactJSON, error) {
	pUrl := fmt.Sprintf("nodes/%s/facts", node)
	ret := []FactJSON{}
	err := c.Get(&ret, pUrl, nil)
	return ret, err
}

func (c *Client) FactPerNode(fact string) ([]FactJSON, error) {
	pUrl := fmt.Sprintf("facts/%s", fact)
	ret := []FactJSON{}
	err := c.Get(&ret, pUrl, nil)
	return ret, err
}

func (c *Client) EventCounts(query string, summarizeBy string, extraParams map[string]string) ([]EventCountJSON, error) {
	path := "event-counts"
	ret := []EventCountJSON{}
	params := mergeParam("query", query, extraParams)
	params = mergeParam("summarize-by", summarizeBy, params)
	err := c.Get(&ret, path, params)
	return ret, err
}

func (c *Client) Events(query string, extraParams map[string]string) ([]EventJSON, error) {
	path := "events"
	ret := []EventJSON{}
	params := mergeParam("query", query, extraParams)
	err := c.Get(&ret, path, params)
	return ret, err
}

//Resources will fetch resources from /resources/ in the puppetdb api
func (c *Client) Resources(query string, extraParams map[string]string) ([]Resource, error) {
	in := []Resource{}
	params := mergeParam("query", query, extraParams)
	err := c.Get(&in, "resources", params)
	return in, err
}

func (c *Client) Metric(v interface{}, metric string) error {
	pUrl := fmt.Sprintf("metrics/mbean/%s", metric)
	err := c.Get(&v, pUrl, nil)
	return err
}

func (c *Client) MetricResourcesPerNode() (result float64, err error) {
	ret := ValueMetricJSON{}
	return ret.Value, c.Metric(&ret, "com.puppetlabs.puppetdb.query.population:type=default,name=avg-resources-per-node")
}

func (c *Client) MetricNumResources() (result float64, err error) {
	ret := ValueMetricJSON{}
	return ret.Value, c.Metric(&ret, "com.puppetlabs.puppetdb.query.population:type=default,name=num-resources")
}

func (c *Client) MetricNumNodes() (result float64, err error) {
	ret := ValueMetricJSON{}
	return ret.Value, c.Metric(&ret, "com.puppetlabs.puppetdb.query.population:type=default,name=num-nodes")
}

func (c *Client) Reports(query string, extraParams map[string]string) ([]ReportJSON, error) {
	path := "reports"
	ret := []ReportJSON{}
	params := mergeParam("query", query, extraParams)
	err := c.Get(&ret, path, params)
	return ret, err
}

func (c *Client) PuppetdbVersion() (PuppetdbVersion, error) {
	path := "version"
	ret := PuppetdbVersion{}
	err := c.Get(&ret, path, nil)
	return ret, err
}

func QueryToJson(query interface{}) (result string, err error) {
	resultBytes, err := json.Marshal(query)
	jsonQuery := string(resultBytes[:])
	return jsonQuery, err
}

func mergeParam(paramName string, paramValue string, params map[string]string) map[string]string {
	resultParams := make(map[string]string)
	if paramValue != "" {
		resultParams[paramName] = paramValue
	}
	if params != nil && len(params) > 0 {
		for k, v := range params {
			resultParams[k] = v
		}
	}
	return resultParams
}

func (c *Client) httpGet(endpoint string) (resp *http.Response, err error) {
	base := strings.TrimRight(c.BaseURL, "/")
	pUrl := fmt.Sprintf("%s/pdb/query/v4/%s", base, endpoint)
	if c.verbose == true {
		log.Printf(pUrl)
	}
	return c.httpClient.Get(pUrl)
}
<commit_msg>Ca certificate needs to included i think from example https://gist.github.com/michaljemala/d6f4e01c4834bf47a9c4<commit_after>package puppetdb

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// Client This represents a connection to your puppetdb instance
type Client struct {
	BaseURL    string
	Cert       string
	Key        string
	httpClient *http.Client
	verbose    bool
}

// EventCountJSON A json object holding the results of a query to the eventcount api
type EventCountJSON struct {
	SubjectType string            `json:"subject-type"`
	Subject     map[string]string `json:"subject"`
	Failure     int64             `json:"failures"`
	Successes   int64             `json:"successes"`
	Noops       int64             `json:"noops"`
	Skips       int64             `json:"skips"`
}

// EventJSON A json object holding the results of a query to the event api.
type EventJSON struct {
	CertName          string `json:"certname"`
	OldValue          string `json:"old-value"`
	Property          string `json:"message"`
	Timestamp         string `json:"timestamp"`
	ResourceType      string `json:"resource-typ"`
	ResourceTitle     string `json:"resource-title"`
	NewValue          string `json:"new-value"`
	Message           string `json:"message"`
	Report            string `json:"report"`
	Status            string `json:"status"`
	File              string `json:"file"`
	ContainmentPath   string `json:"containment-path"`
	ContainmentClass  string `json:"containing-class"`
	RunStartTime      string `json:"run-start-time"`
	RunEndTime        string `json:"run-end-time"`
	ReportReceiveTime string `json:"report-receive-time"`
}

// FactJSON A json object holding the results of a query to the facts api.
type FactJSON struct {
	CertName    string `json:"certname"`
	Environment string `json:"environment"`
	Name        string `json:"name"`
	Value       string `json:"value"`
}

// NodeJSON A json object holding the results of query to the node api.
type NodeJSON struct {
	Certname                     string `json:"certname"`
	Deactivated                  string `json:"deactivated"`
	CatalogTimestamp             string `json:"catalog_timestamp"`
	FactsTimestamp               string `json:"facts_timestamp"`
	CatalogEnvironment           string `json:"catalog_environment"`
	FactsEnvironment             string `json:"facts_environment"`
	LatestReportHash             string `json:"latest_report_hash"`
	CachedCatalogStatus          string `json:"cached_catalog_status"`
	ReportEnvironment            string `json:"report_environment"`
	ReportTimestamp              string `json:"report_timestamp"`
	LatestReportCorrectiveChange string `json:"latest_report_corrective_change"`
	LatestReportNoop             bool   `json:"latest_report_noop"`
	LatestReportNoopPending      bool   `json:"latest_report_noop_pending"`
	Expired                      string `json:"expired"`
	LatestReportJobID            string `json:"latest_report_job_id"`
	LatestReportStatus           string `json:"latest_report_status"`
}

// PuppetdbVersion a simple struct holding the puppetdb version.
type PuppetdbVersion struct {
	Version string `json:"version"`
}

// ReportJSON A json abject holding the data for a query from the report api.
type ReportJSON struct {
	CertName             string `json:"certname"`
	PuppetVersion        string `json:"puppet-version"`
	Value                string `json:"value"`
	Hash                 string `json:"hash"`
	ReportFormat         int64  `json:"report-format"`
	ConfigurationVersion string `json:"configuration-version"`
	TransactionUUID      string `json:"transaction-uuid"`
	StartTime            string `json:"start-time"`
	EndTime              string `json:"end-time"`
	ReceiveTime          string `json:"receive-time"`
}

//Resource contains information about a puppet resource.
type Resource struct {
	Paramaters map[string]interface{} `json:"parameters"`
	Line       int                    `json:"line,omitempty"`
	Exported   bool                   `json:"exported,omitempty"`
	Tags       []string               `json:"tags,omitempty"`
	Title      string                 `json:"title,omitempty"`
	Type       string                 `json:"type,omitempty"`
	Resource   string                 `json:"resource,omitempty"`
	Certname   string                 `json:"certname,omitempty"`
}

// ValueMetricJSON A simple structholding a float value.
type ValueMetricJSON struct {
	Value float64
}

// getURL return the address of the puppetdb instance.
func getURL(host string, port int, ssl bool) string {
	if ssl {
		return fmt.Sprintf("https://%s:%v", host, port)
	} else {
		return fmt.Sprintf("http://%s:%v", host, port)
	}
}

// NewClient returns a http connection for your puppetdb instance.
func NewClient(host string, port int, verbose bool) *Client {

	tr := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	client := &http.Client{Transport: tr}
	return &Client{getURL(host, port, false), "", "", client, verbose}
}

// NewClientSSL returns a https connection for your puppetdb instance.
func NewClientSSL(host string, port int, key string, cert string, ca string, verbose bool) *Client {
	flag.Parse()
	cert2, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		log.Fatal(err)
	}
	// Load CA cert
	caCert, err := ioutil.ReadFile(ca)
	if err != nil {
		log.Fatal(err)
	}
	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caCert)

	// Setup HTTPS client
	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert2},
		RootCAs:      caCertPool,
	}
	tlsConfig.BuildNameToCertificate()
	transport := &http.Transport{TLSClientConfig: tlsConfig}
	client := &http.Client{Transport: transport}
	return &Client{getURL(host, port, true), cert, key, client, verbose}

}

// NewClientTimeout returns a http connection for your puppetdb instance with a timeout.
func NewClientTimeout(host string, port int, verbose bool, timeout int) *Client {

	tr := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	client := &http.Client{Transport: tr, Timeout: time.Duration(timeout) * time.Second}
	return &Client{getURL(host, port, false), "", "", client, verbose}
}

// NewClientTimeoutSSL returns a http connection for your puppetdb instance with a timeout and ssl configured.
func NewClientTimeoutSSL(host string, port int, key string, cert string, ca string, verbose bool, timeout int) *Client {
	flag.Parse()
	cert2, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		log.Fatal(err)
	}
	// Load CA cert
	caCert, err := ioutil.ReadFile(ca)
	if err != nil {
		log.Fatal(err)
	}
	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caCert)

	// Setup HTTPS client
	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert2},
		RootCAs:      caCertPool,
	}
	tlsConfig.BuildNameToCertificate()
	transport := &http.Transport{TLSClientConfig: tlsConfig}
	client := &http.Client{Transport: transport, Timeout: time.Duration(timeout) * time.Second}
	return &Client{getURL(host, port, true), cert, key, client, verbose}

}

// Get gets the given url and retrusn the result. In form of the given interface.
func (c *Client) Get(v interface{}, path string, params map[string]string) error {
	pathAndParams := path
	//TODO: Improve this
	if params != nil && len(params) > 0 {
		if !strings.Contains(path, "?") {
			pathAndParams += "?"
		}
		for k, v := range params {
			pathAndParams += fmt.Sprintf("%s=%s&", k, url.QueryEscape(v))
		}
	}
	resp, err := c.httpGet(pathAndParams)
	if err != nil {
		log.Print(err)
		return err
	}
	defer resp.Body.Close()
	if err != nil {
		log.Print(err)
		return err
	}
	json.NewDecoder(resp.Body).Decode(&v)
	return err
}

// Nodes Polls the nodes api of your puppetdb and returns the results in form of the NodeJSON type.
func (c *Client) Nodes() ([]NodeJSON, error) {
	ret := []NodeJSON{}
	err := c.Get(&ret, "nodes", nil)
	return ret, err
}

func (c *Client) FactNames() ([]string, error) {
	ret := []string{}
	err := c.Get(&ret, "fact-names", nil)
	return ret, err
}

func (c *Client) NodeFacts(node string) ([]FactJSON, error) {
	pUrl := fmt.Sprintf("nodes/%s/facts", node)
	ret := []FactJSON{}
	err := c.Get(&ret, pUrl, nil)
	return ret, err
}

func (c *Client) FactPerNode(fact string) ([]FactJSON, error) {
	pUrl := fmt.Sprintf("facts/%s", fact)
	ret := []FactJSON{}
	err := c.Get(&ret, pUrl, nil)
	return ret, err
}

func (c *Client) EventCounts(query string, summarizeBy string, extraParams map[string]string) ([]EventCountJSON, error) {
	path := "event-counts"
	ret := []EventCountJSON{}
	params := mergeParam("query", query, extraParams)
	params = mergeParam("summarize-by", summarizeBy, params)
	err := c.Get(&ret, path, params)
	return ret, err
}

func (c *Client) Events(query string, extraParams map[string]string) ([]EventJSON, error) {
	path := "events"
	ret := []EventJSON{}
	params := mergeParam("query", query, extraParams)
	err := c.Get(&ret, path, params)
	return ret, err
}

//Resources will fetch resources from /resources/ in the puppetdb api
func (c *Client) Resources(query string, extraParams map[string]string) ([]Resource, error) {
	in := []Resource{}
	params := mergeParam("query", query, extraParams)
	err := c.Get(&in, "resources", params)
	return in, err
}

func (c *Client) Metric(v interface{}, metric string) error {
	pUrl := fmt.Sprintf("metrics/mbean/%s", metric)
	err := c.Get(&v, pUrl, nil)
	return err
}

func (c *Client) MetricResourcesPerNode() (result float64, err error) {
	ret := ValueMetricJSON{}
	return ret.Value, c.Metric(&ret, "com.puppetlabs.puppetdb.query.population:type=default,name=avg-resources-per-node")
}

func (c *Client) MetricNumResources() (result float64, err error) {
	ret := ValueMetricJSON{}
	return ret.Value, c.Metric(&ret, "com.puppetlabs.puppetdb.query.population:type=default,name=num-resources")
}

func (c *Client) MetricNumNodes() (result float64, err error) {
	ret := ValueMetricJSON{}
	return ret.Value, c.Metric(&ret, "com.puppetlabs.puppetdb.query.population:type=default,name=num-nodes")
}

func (c *Client) Reports(query string, extraParams map[string]string) ([]ReportJSON, error) {
	path := "reports"
	ret := []ReportJSON{}
	params := mergeParam("query", query, extraParams)
	err := c.Get(&ret, path, params)
	return ret, err
}

func (c *Client) PuppetdbVersion() (PuppetdbVersion, error) {
	path := "version"
	ret := PuppetdbVersion{}
	err := c.Get(&ret, path, nil)
	return ret, err
}

func QueryToJson(query interface{}) (result string, err error) {
	resultBytes, err := json.Marshal(query)
	jsonQuery := string(resultBytes[:])
	return jsonQuery, err
}

func mergeParam(paramName string, paramValue string, params map[string]string) map[string]string {
	resultParams := make(map[string]string)
	if paramValue != "" {
		resultParams[paramName] = paramValue
	}
	if params != nil && len(params) > 0 {
		for k, v := range params {
			resultParams[k] = v
		}
	}
	return resultParams
}

func (c *Client) httpGet(endpoint string) (resp *http.Response, err error) {
	base := strings.TrimRight(c.BaseURL, "/")
	pUrl := fmt.Sprintf("%s/pdb/query/v4/%s", base, endpoint)
	if c.verbose == true {
		log.Printf(pUrl)
	}
	return c.httpClient.Get(pUrl)
}
<|endoftext|>


<commit_before>// Copyright 2018 The Terraformer Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcp

import (
	"context"
	"log"
	"regexp"

	"cloud.google.com/go/iam/admin/apiv1"
	"google.golang.org/api/cloudresourcemanager/v1"
	"google.golang.org/api/iterator"
	adminpb "google.golang.org/genproto/googleapis/iam/admin/v1"

	"github.com/GoogleCloudPlatform/terraformer/terraform_utils"
)

var IamAllowEmptyValues = []string{"tags."}

var IamAdditionalFields = map[string]interface{}{}

type IamGenerator struct {
	GCPService
}

func (IamGenerator) createServiceAccountResources(serviceAccountsIterator *admin.ServiceAccountIterator) []terraform_utils.Resource {
	resources := []terraform_utils.Resource{}
	re := regexp.MustCompile(`^[a-z]`)
	for {
		serviceAccount, err := serviceAccountsIterator.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			log.Println("error with service account:", err)
			continue
		}
		if !re.MatchString(serviceAccount.Email) {
			log.Printf("skipping %s: service account email must start with [a-z]\n", serviceAccount.Name)
			continue
		}
		resources = append(resources, terraform_utils.NewSimpleResource(
			serviceAccount.Name,
			serviceAccount.UniqueId,
			"google_service_account",
			"google",
			IamAllowEmptyValues,
		))
	}
	return resources
}

func (g *IamGenerator) createIamCustomRoleResources(rolesResponse *adminpb.ListRolesResponse, project string) []terraform_utils.Resource {
	resources := []terraform_utils.Resource{}
	for _, role := range rolesResponse.Roles {
		resources = append(resources, terraform_utils.NewResource(
			role.Name,
			role.Name,
			"google_project_iam_custom_role",
			"google",
			map[string]string{
				"role_id": role.Name,
				"project": project,
			},
			IamAllowEmptyValues,
			map[string]interface{}{
				"stage": role.Stage.String(),
			},
		))
	}

	return resources
}

func (g *IamGenerator) createIamBindingResources(policy *cloudresourcemanager.Policy, project string) []terraform_utils.Resource {
	resources := []terraform_utils.Resource{}
	for _, b := range policy.Bindings {
		resources = append(resources, terraform_utils.NewResource(
			b.Role,
			b.Role,
			"google_project_iam_binding",
			"google",
			map[string]string{
				"role":    b.Role,
				"project": project,
			},
			IamAllowEmptyValues,
			IamAdditionalFields,
		))
	}

	return resources
}

func (g *IamGenerator) InitResources() error {
	ctx := context.Background()

	projectID := g.GetArgs()["project"].(string)
	client, err := admin.NewIamClient(ctx)
	if err != nil {
		log.Fatal(err)
	}
	serviceAccountsIterator := client.ListServiceAccounts(ctx, &adminpb.ListServiceAccountsRequest{Name: "projects/" + projectID})
	rolesResponse, err := client.ListRoles(ctx, &adminpb.ListRolesRequest{Parent: "projects/" + projectID})
	if err != nil {
		log.Fatal(err)
	}

	cm, err := cloudresourcemanager.NewService(context.Background())
	if err != nil {
		log.Fatal(err)
	}
	rb := &cloudresourcemanager.GetIamPolicyRequest{}
	policyResponse, err := cm.Projects.GetIamPolicy(projectID, rb).Context(context.Background()).Do()
	if err != nil {
		log.Fatal(err)
	}

	g.Resources = g.createServiceAccountResources(serviceAccountsIterator)
	g.Resources = append(g.Resources, g.createIamCustomRoleResources(rolesResponse, projectID)...)
	g.Resources = append(g.Resources, g.createIamBindingResources(policyResponse, projectID)...)
	return nil

}
<commit_msg>Skip IAM Custom Role when it is marked as Deleted<commit_after>// Copyright 2018 The Terraformer Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcp

import (
	"context"
	"log"
	"regexp"

	"cloud.google.com/go/iam/admin/apiv1"
	"google.golang.org/api/cloudresourcemanager/v1"
	"google.golang.org/api/iterator"
	adminpb "google.golang.org/genproto/googleapis/iam/admin/v1"

	"github.com/GoogleCloudPlatform/terraformer/terraform_utils"
)

var IamAllowEmptyValues = []string{"tags."}

var IamAdditionalFields = map[string]interface{}{}

type IamGenerator struct {
	GCPService
}

func (IamGenerator) createServiceAccountResources(serviceAccountsIterator *admin.ServiceAccountIterator) []terraform_utils.Resource {
	resources := []terraform_utils.Resource{}
	re := regexp.MustCompile(`^[a-z]`)
	for {
		serviceAccount, err := serviceAccountsIterator.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			log.Println("error with service account:", err)
			continue
		}
		if !re.MatchString(serviceAccount.Email) {
			log.Printf("skipping %s: service account email must start with [a-z]\n", serviceAccount.Name)
			continue
		}
		resources = append(resources, terraform_utils.NewSimpleResource(
			serviceAccount.Name,
			serviceAccount.UniqueId,
			"google_service_account",
			"google",
			IamAllowEmptyValues,
		))
	}
	return resources
}

func (g *IamGenerator) createIamCustomRoleResources(rolesResponse *adminpb.ListRolesResponse, project string) []terraform_utils.Resource {
	resources := []terraform_utils.Resource{}
	for _, role := range rolesResponse.Roles {
		if role.Deleted {
			// Note: no need to log that the resource has been deleted
			continue
		}
		resources = append(resources, terraform_utils.NewResource(
			role.Name,
			role.Name,
			"google_project_iam_custom_role",
			"google",
			map[string]string{
				"role_id": role.Name,
				"project": project,
			},
			IamAllowEmptyValues,
			map[string]interface{}{
				"stage": role.Stage.String(),
			},
		))
	}

	return resources
}

func (g *IamGenerator) createIamBindingResources(policy *cloudresourcemanager.Policy, project string) []terraform_utils.Resource {
	resources := []terraform_utils.Resource{}
	for _, b := range policy.Bindings {
		resources = append(resources, terraform_utils.NewResource(
			b.Role,
			b.Role,
			"google_project_iam_binding",
			"google",
			map[string]string{
				"role":    b.Role,
				"project": project,
			},
			IamAllowEmptyValues,
			IamAdditionalFields,
		))
	}

	return resources
}

func (g *IamGenerator) InitResources() error {
	ctx := context.Background()

	projectID := g.GetArgs()["project"].(string)
	client, err := admin.NewIamClient(ctx)
	if err != nil {
		log.Fatal(err)
	}
	serviceAccountsIterator := client.ListServiceAccounts(ctx, &adminpb.ListServiceAccountsRequest{Name: "projects/" + projectID})
	rolesResponse, err := client.ListRoles(ctx, &adminpb.ListRolesRequest{Parent: "projects/" + projectID})
	if err != nil {
		log.Fatal(err)
	}

	cm, err := cloudresourcemanager.NewService(context.Background())
	if err != nil {
		log.Fatal(err)
	}
	rb := &cloudresourcemanager.GetIamPolicyRequest{}
	policyResponse, err := cm.Projects.GetIamPolicy(projectID, rb).Context(context.Background()).Do()
	if err != nil {
		log.Fatal(err)
	}

	g.Resources = g.createServiceAccountResources(serviceAccountsIterator)
	g.Resources = append(g.Resources, g.createIamCustomRoleResources(rolesResponse, projectID)...)
	g.Resources = append(g.Resources, g.createIamBindingResources(policyResponse, projectID)...)
	return nil

}
<|endoftext|>


<commit_before>package http

import (
	"context"
	"encoding/base64"
	"io"
	"strings"

	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/retry"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/common/task"
	"v2ray.com/core/features/policy"
	"v2ray.com/core/transport"
	"v2ray.com/core/transport/internet"
)

type Client struct {
	serverPicker  protocol.ServerPicker
	policyManager policy.Manager
}

// NewClient create a new http client based on the given config.
func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
	serverList := protocol.NewServerList()
	for _, rec := range config.Server {
		s, err := protocol.NewServerSpecFromPB(*rec)
		if err != nil {
			return nil, newError("failed to get server spec").Base(err)
		}
		serverList.AddServer(s)
	}
	if serverList.Size() == 0 {
		return nil, newError("0 target server")
	}

	v := core.MustFromContext(ctx)
	return &Client{
		serverPicker:  protocol.NewRoundRobinServerPicker(serverList),
		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
	}, nil
}

// Process implements proxy.Outbound.Process.
// 使用connect方法连接http代理服务器，获得一个隧道，然后通过该隧道通信
func (c *Client) Process(ctx context.Context, link *transport.Link, dialer internet.Dialer) error {
	outbound := session.OutboundFromContext(ctx)
	if outbound == nil || !outbound.Target.IsValid() {
		return newError("target not specified.")
	}
	destination := outbound.Target

	if destination.Network == net.Network_UDP {
		return newError("UDP is not supported by HTTP outbound")
	}

	var server *protocol.ServerSpec
	var conn internet.Connection

	if err := retry.ExponentialBackoff(5, 100).On(func() error {
		server = c.serverPicker.PickServer()
		dest := server.Destination()
		rawConn, err := dialer.Dial(ctx, dest)
		if err != nil {
			return err
		}
		conn = rawConn

		return nil
	}); err != nil {
		return newError("failed to find an available destination").Base(err)
	}

	defer func() {
		if err := conn.Close(); err != nil {
			newError("failed to closed connection").Base(err).WriteToLog(session.ExportIDToError(ctx))
		}
	}()

	p := c.policyManager.ForLevel(0)

	user := server.PickUser()
	if user != nil {
		p = c.policyManager.ForLevel(user.Level)
	}

	if err := setUpHttpTunnel(conn, conn, &destination, user); err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(ctx)
	timer := signal.CancelAfterInactivity(ctx, cancel, p.Timeouts.ConnectionIdle)

	requestFunc := func() error {
		defer timer.SetTimeout(p.Timeouts.DownlinkOnly)
		return buf.Copy(link.Reader, buf.NewWriter(conn), buf.UpdateActivity(timer))
	}
	responseFunc := func() error {
		defer timer.SetTimeout(p.Timeouts.UplinkOnly)
		return buf.Copy(buf.NewReader(conn), link.Writer, buf.UpdateActivity(timer))
	}

	var responseDonePost = task.OnSuccess(responseFunc, task.Close(link.Writer))
	if err := task.Run(ctx, requestFunc, responseDonePost); err != nil {
		return newError("connection ends").Base(err)
	}

	return nil
}

// 使用http connect方法建立一个隧道
func setUpHttpTunnel(reader io.Reader, writer io.Writer, destination *net.Destination, user *protocol.MemoryUser) error {
	var headers []string
	destNetAddr := destination.NetAddr()
	headers = append(headers, "CONNECT "+destNetAddr+" HTTP/1.1")
	headers = append(headers, "Host: "+destNetAddr)
	if user != nil && user.Account != nil {
		account := user.Account.(*Account)
		auth := account.GetUsername() + ":" + account.GetPassword()
		headers = append(headers, "Proxy-Authorization: Basic "+base64.StdEncoding.EncodeToString([]byte(auth)))
	}
	headers = append(headers, "Proxy-Connection: Keep-Alive")

	b := buf.New()
	b.WriteString(strings.Join(headers, "\r\n") + "\r\n\r\n")
	if err := buf.WriteAllBytes(writer, b.Bytes()); err != nil {
		return err
	}

	b.Clear()
	if _, err := b.ReadFrom(reader); err != nil {
		return err
	}

	return nil
}

func init() {
	common.Must(common.RegisterConfig((*ClientConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewClient(ctx, config.(*ClientConfig))
	}))
}
<commit_msg>Rewrite comments in English<commit_after>package http

import (
	"context"
	"encoding/base64"
	"io"
	"strings"

	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/retry"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/common/task"
	"v2ray.com/core/features/policy"
	"v2ray.com/core/transport"
	"v2ray.com/core/transport/internet"
)

type Client struct {
	serverPicker  protocol.ServerPicker
	policyManager policy.Manager
}

// NewClient create a new http client based on the given config.
func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
	serverList := protocol.NewServerList()
	for _, rec := range config.Server {
		s, err := protocol.NewServerSpecFromPB(*rec)
		if err != nil {
			return nil, newError("failed to get server spec").Base(err)
		}
		serverList.AddServer(s)
	}
	if serverList.Size() == 0 {
		return nil, newError("0 target server")
	}

	v := core.MustFromContext(ctx)
	return &Client{
		serverPicker:  protocol.NewRoundRobinServerPicker(serverList),
		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
	}, nil
}

// Process implements proxy.Outbound.Process. We first create a socket tunnel via HTTP CONNECT method, then redirect all inbound traffic to that tunnel.
func (c *Client) Process(ctx context.Context, link *transport.Link, dialer internet.Dialer) error {
	outbound := session.OutboundFromContext(ctx)
	if outbound == nil || !outbound.Target.IsValid() {
		return newError("target not specified.")
	}
	destination := outbound.Target

	if destination.Network == net.Network_UDP {
		return newError("UDP is not supported by HTTP outbound")
	}

	var server *protocol.ServerSpec
	var conn internet.Connection

	if err := retry.ExponentialBackoff(5, 100).On(func() error {
		server = c.serverPicker.PickServer()
		dest := server.Destination()
		rawConn, err := dialer.Dial(ctx, dest)
		if err != nil {
			return err
		}
		conn = rawConn

		return nil
	}); err != nil {
		return newError("failed to find an available destination").Base(err)
	}

	defer func() {
		if err := conn.Close(); err != nil {
			newError("failed to closed connection").Base(err).WriteToLog(session.ExportIDToError(ctx))
		}
	}()

	p := c.policyManager.ForLevel(0)

	user := server.PickUser()
	if user != nil {
		p = c.policyManager.ForLevel(user.Level)
	}

	if err := setUpHttpTunnel(conn, conn, &destination, user); err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(ctx)
	timer := signal.CancelAfterInactivity(ctx, cancel, p.Timeouts.ConnectionIdle)

	requestFunc := func() error {
		defer timer.SetTimeout(p.Timeouts.DownlinkOnly)
		return buf.Copy(link.Reader, buf.NewWriter(conn), buf.UpdateActivity(timer))
	}
	responseFunc := func() error {
		defer timer.SetTimeout(p.Timeouts.UplinkOnly)
		return buf.Copy(buf.NewReader(conn), link.Writer, buf.UpdateActivity(timer))
	}

	var responseDonePost = task.OnSuccess(responseFunc, task.Close(link.Writer))
	if err := task.Run(ctx, requestFunc, responseDonePost); err != nil {
		return newError("connection ends").Base(err)
	}

	return nil
}

// setUpHttpTunnel will create a socket tunnel via HTTP CONNECT method
func setUpHttpTunnel(reader io.Reader, writer io.Writer, destination *net.Destination, user *protocol.MemoryUser) error {
	var headers []string
	destNetAddr := destination.NetAddr()
	headers = append(headers, "CONNECT "+destNetAddr+" HTTP/1.1")
	headers = append(headers, "Host: "+destNetAddr)
	if user != nil && user.Account != nil {
		account := user.Account.(*Account)
		auth := account.GetUsername() + ":" + account.GetPassword()
		headers = append(headers, "Proxy-Authorization: Basic "+base64.StdEncoding.EncodeToString([]byte(auth)))
	}
	headers = append(headers, "Proxy-Connection: Keep-Alive")

	b := buf.New()
	b.WriteString(strings.Join(headers, "\r\n") + "\r\n\r\n")
	if err := buf.WriteAllBytes(writer, b.Bytes()); err != nil {
		return err
	}

	b.Clear()
	if _, err := b.ReadFrom(reader); err != nil {
		return err
	}

	return nil
}

func init() {
	common.Must(common.RegisterConfig((*ClientConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewClient(ctx, config.(*ClientConfig))
	}))
}
<|endoftext|>


<commit_before>package pop3

import (
	"bds/lib/mailstore"
	"bds/lib/starttls"
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"io"
	"net"
	"net/textproto"
	"os"
	"strconv"
	"strings"
)

// function that authenticates a user
type UserAuthenticator func(string, string) (bool, error)

// pop3 server
type Server struct {
	// obtains a mail store given a user
	Local mailstore.MailRouter
	// login authenticator
	Auth UserAuthenticator
	// server name
	name string
	// tls config
	TLS *tls.Config
}

func New() *Server {
	host, _ := os.Hostname()
	return &Server{
		name: host,
	}
}

// get all messages in mail store
func (s *Server) getMessages(user string) (msgs []mailstore.Message, err error) {
	// get user's mail store
	var st mailstore.Store
	if s.Local == nil {
		err = errors.New("could't find mail store")
		return
	}
	var has bool
	st, has = s.Local.FindStoreFor(user)
	if !has {
		err = errors.New("no such local user")
		return
	}
	var ms []mailstore.Message
	ms, err = st.ListNew()
	// move new mail into cur	msgs
	if err == nil {
		for _, m := range ms {
			_, err = st.Process(m)
			if err != nil {
				log.Errorf("error processing maildir: %s", err.Error())
			}
		}
	}
	// get list of messages
	ms, err = st.List()
	if err == nil {
		for _, msg := range ms {
			msgs = append(msgs, msg)
		}
	}
	return
}

func (s *Server) checkUser(user, passwd string) (allowed bool) {
	if s.Auth != nil {
		allowed, _ = s.Auth(user, passwd)
	}
	return
}

// get all messages and octet count
func (s *Server) obtainMessages(user string) (msgs []mailstore.Message, o int64, err error) {
	msgs, err = s.getMessages(user)
	if err == nil {
		for _, msg := range msgs {
			var info os.FileInfo
			info, err = os.Stat(msg.Filepath())
			if err == nil && !info.IsDir() {
				o += info.Size()
			}
		}
	} else {
		log.Errorf("pop3: %s", err.Error())
	}
	return
}

// pop3 session handler
type pop3Session struct {
	// network connection
	nc net.Conn
	c  *textproto.Conn
	// parent server
	s *Server
	// are we in transaction state?
	transaction bool
	// current user
	user string
	// messages we have in this transaction
	msgs []mailstore.Message
	// how many octets we have for all messages
	octs int64
	// messages to delete
	dels []mailstore.Message
}

// run pop3 session mainloop
func (p *pop3Session) Run() {
	// send banner
	err := p.OK("POP3 Server Ready")
	for err == nil {
		var line string
		line, err = p.c.ReadLine()
		if err == nil {
			if strings.ToUpper(line) == "QUIT" {
				// check for quit command
				err = p.OK("k bai")
				break
			} else if p.transaction {
				err = p.handleTransactionLine(line)
			} else {
				err = p.handleLine(line)
			}
		}
	}
	if err != nil && err != io.EOF {
		log.Errorf("error in pop3 session: %s", err.Error())
	}
	// close connection
	p.c.Close()
	// delete old messages
	for _, msg := range p.dels {
		os.Remove(msg.Filepath())
	}
}

// handle line when in transaction mode
func (p *pop3Session) handleTransactionLine(line string) (err error) {
	var info os.FileInfo
	var idx int
	parts := strings.Split(line, " ")
	cmd := strings.ToUpper(parts[0])
	switch cmd {
	case "DELE":
		if len(parts) == 2 {
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid, add it to delete
				p.dels = append(p.dels, p.msgs[idx-1])
				p.OK("")
			} else {
				// invalid
				err = p.Error(err.Error())
			}
		}
	case "RETR":
		if len(parts) == 2 {
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid
				msg := p.msgs[idx-1].Filepath()
				var f *os.File
				f, err = os.Open(msg)
				if err == nil {
					r := bufio.NewReader(f)
					info, err = f.Stat()
					if err == nil {
						err = p.c.PrintfLine("+OK %d octets", info.Size())
						for err == nil {
							// send line
							line, err = r.ReadString(10)
							if err == io.EOF {
								err = nil
								break
							} else if err == nil {
								line = strings.Trim(line, "\r")
								line = strings.Trim(line, "\n")
								if line == "." {
									line = " ."
								}
								// send line
								err = p.c.PrintfLine(line)
							} else {
								// error
								break
							}
						}
					}
					f.Close()
					// end
					err = p.c.PrintfLine(".")
				} else {
					err = p.Error(err.Error())
				}
			} else {
				// invalid
				err = p.Error("bad message")
			}
		} else {
			err = p.Error("invalid syntax")
		}
		break
	case "UIDL":
		if len(parts) == 2 {
			// 1 message
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid
				err = p.OK(parts[1] + " " + p.msgs[idx-1].Filename())
			} else {
				// invalid
				err = p.Error("bad message")
			}
		} else {
			// all messages
			p.OK("")
			dw := p.c.DotWriter()
			for idx, msg := range p.msgs {
				fmt.Fprintf(dw, "%d %s\r\n", 1+idx, msg.Filename())
			}
			// FLUSH :D
			err = dw.Close()
		}
		// begin
		break
	case "STAT":
		// begin
		_, err = fmt.Fprintf(p.c.W, "+OK %d %d", len(p.msgs), p.octs)

		if err == nil {
			// done writing list
			err = p.c.PrintfLine("")
		} else {
			err = p.Error(err.Error())
		}
		break
	case "LIST":
		if len(parts) == 2 {
			// 1 message
			idx, err = strconv.Atoi(parts[1])
			if err == nil {
				if idx > 0 && idx <= len(p.msgs) {
					info, err = os.Stat(p.msgs[idx-1].Filepath())
					if err == nil {
						err = p.c.PrintfLine("+OK %d %d", idx, info.Size())
					}
				} else {
					// no existing message
					err = p.Error("no such message")
				}
			}
		} else {
			// all messages
			err = p.c.PrintfLine("+OK %d messages (%d octets)", len(p.msgs), p.octs)
			if err == nil {
				dw := p.c.DotWriter()
				for i, msg := range p.msgs {
					info, err = os.Stat(msg.Filepath())
					if err == nil {
						// write out entry
						fmt.Fprintf(dw, "%d %d\r\n", i+1, info.Size())
					} else {
						log.Errorf("error in pop3, stat(): %s", err.Error())
					}
				}
				// FLUSH IT !!! :-DDDDD
				err = dw.Close()
			}
		}
		break
	default:
		err = p.Error("bad command")
	}
	return
}

// handle 1 line of input when not in transaction mode
func (p *pop3Session) handleLine(line string) (err error) {
	parts := strings.Split(line, " ")
	cmd := strings.ToUpper(parts[0])
	c := p.c
	switch cmd {
	case "STARTTLS":
		if p.s.TLS == nil {
			err = p.Error("no TLS supported")
		} else {
			p.OK("Begin TLS negotiation")
			p.c, _, err = starttls.HandleStartTLS(p.nc, p.s.TLS)
			if err != nil {
				c.Close()
			}
		}
	case "NOOP":
		err = p.OK("")
		break
	case "PASS":
		if p.s.checkUser(p.user, line[5:]) {
			p.msgs, p.octs, err = p.s.obtainMessages(p.user)
			if err == nil {
				err = p.c.PrintfLine("+OK %s maildrop logged in, you have %d messages (%d octets)", p.user, len(p.msgs), p.octs)
				p.transaction = err == nil
			} else {
				err = p.Error(err.Error())
			}
		} else {
			err = p.Error("bad login")
		}
		break
	case "USER":
		if len(parts) > 1 {
			p.user = line[5:]
		}
		err = p.OK("you may try login as " + p.user)
		break
	default:
		err = p.Error("bad command")
	}
	return
}

func (p *pop3Session) OK(msg string) (err error) {
	if len(msg) == 0 {
		err = p.c.PrintfLine("+OK")
	} else {
		err = p.c.PrintfLine("+OK %s", msg)
	}
	return
}

// send error
func (p *pop3Session) Error(msg string) (err error) {
	err = p.c.PrintfLine("-ERR %s", msg)
	return
}

// serve sessions with connections accepted from a net.Listener
func (s *Server) Serve(l net.Listener) (err error) {
	for err == nil {
		var c net.Conn
		c, err = l.Accept()
		if err == nil {
			p := &pop3Session{
				nc: c,
				c:  textproto.NewConn(c),
				s:  s,
			}
			go p.Run()
		}
	}
	return
}
<commit_msg>starttls -> stls<commit_after>package pop3

import (
	"bds/lib/mailstore"
	"bds/lib/starttls"
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"io"
	"net"
	"net/textproto"
	"os"
	"strconv"
	"strings"
)

// function that authenticates a user
type UserAuthenticator func(string, string) (bool, error)

// pop3 server
type Server struct {
	// obtains a mail store given a user
	Local mailstore.MailRouter
	// login authenticator
	Auth UserAuthenticator
	// server name
	name string
	// tls config
	TLS *tls.Config
}

func New() *Server {
	host, _ := os.Hostname()
	return &Server{
		name: host,
	}
}

// get all messages in mail store
func (s *Server) getMessages(user string) (msgs []mailstore.Message, err error) {
	// get user's mail store
	var st mailstore.Store
	if s.Local == nil {
		err = errors.New("could't find mail store")
		return
	}
	var has bool
	st, has = s.Local.FindStoreFor(user)
	if !has {
		err = errors.New("no such local user")
		return
	}
	var ms []mailstore.Message
	ms, err = st.ListNew()
	// move new mail into cur	msgs
	if err == nil {
		for _, m := range ms {
			_, err = st.Process(m)
			if err != nil {
				log.Errorf("error processing maildir: %s", err.Error())
			}
		}
	}
	// get list of messages
	ms, err = st.List()
	if err == nil {
		for _, msg := range ms {
			msgs = append(msgs, msg)
		}
	}
	return
}

func (s *Server) checkUser(user, passwd string) (allowed bool) {
	if s.Auth != nil {
		allowed, _ = s.Auth(user, passwd)
	}
	return
}

// get all messages and octet count
func (s *Server) obtainMessages(user string) (msgs []mailstore.Message, o int64, err error) {
	msgs, err = s.getMessages(user)
	if err == nil {
		for _, msg := range msgs {
			var info os.FileInfo
			info, err = os.Stat(msg.Filepath())
			if err == nil && !info.IsDir() {
				o += info.Size()
			}
		}
	} else {
		log.Errorf("pop3: %s", err.Error())
	}
	return
}

// pop3 session handler
type pop3Session struct {
	// network connection
	nc net.Conn
	c  *textproto.Conn
	// parent server
	s *Server
	// are we in transaction state?
	transaction bool
	// current user
	user string
	// messages we have in this transaction
	msgs []mailstore.Message
	// how many octets we have for all messages
	octs int64
	// messages to delete
	dels []mailstore.Message
}

// run pop3 session mainloop
func (p *pop3Session) Run() {
	// send banner
	err := p.OK("POP3 Server Ready")
	for err == nil {
		var line string
		line, err = p.c.ReadLine()
		if err == nil {
			if strings.ToUpper(line) == "QUIT" {
				// check for quit command
				err = p.OK("k bai")
				break
			} else if p.transaction {
				err = p.handleTransactionLine(line)
			} else {
				err = p.handleLine(line)
			}
		}
	}
	if err != nil && err != io.EOF {
		log.Errorf("error in pop3 session: %s", err.Error())
	}
	// close connection
	p.c.Close()
	// delete old messages
	for _, msg := range p.dels {
		os.Remove(msg.Filepath())
	}
}

// handle line when in transaction mode
func (p *pop3Session) handleTransactionLine(line string) (err error) {
	var info os.FileInfo
	var idx int
	parts := strings.Split(line, " ")
	cmd := strings.ToUpper(parts[0])
	switch cmd {
	case "DELE":
		if len(parts) == 2 {
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid, add it to delete
				p.dels = append(p.dels, p.msgs[idx-1])
				p.OK("")
			} else {
				// invalid
				err = p.Error(err.Error())
			}
		}
	case "RETR":
		if len(parts) == 2 {
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid
				msg := p.msgs[idx-1].Filepath()
				var f *os.File
				f, err = os.Open(msg)
				if err == nil {
					r := bufio.NewReader(f)
					info, err = f.Stat()
					if err == nil {
						err = p.c.PrintfLine("+OK %d octets", info.Size())
						for err == nil {
							// send line
							line, err = r.ReadString(10)
							if err == io.EOF {
								err = nil
								break
							} else if err == nil {
								line = strings.Trim(line, "\r")
								line = strings.Trim(line, "\n")
								if line == "." {
									line = " ."
								}
								// send line
								err = p.c.PrintfLine(line)
							} else {
								// error
								break
							}
						}
					}
					f.Close()
					// end
					err = p.c.PrintfLine(".")
				} else {
					err = p.Error(err.Error())
				}
			} else {
				// invalid
				err = p.Error("bad message")
			}
		} else {
			err = p.Error("invalid syntax")
		}
		break
	case "UIDL":
		if len(parts) == 2 {
			// 1 message
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid
				err = p.OK(parts[1] + " " + p.msgs[idx-1].Filename())
			} else {
				// invalid
				err = p.Error("bad message")
			}
		} else {
			// all messages
			p.OK("")
			dw := p.c.DotWriter()
			for idx, msg := range p.msgs {
				fmt.Fprintf(dw, "%d %s\r\n", 1+idx, msg.Filename())
			}
			// FLUSH :D
			err = dw.Close()
		}
		// begin
		break
	case "STAT":
		// begin
		_, err = fmt.Fprintf(p.c.W, "+OK %d %d", len(p.msgs), p.octs)

		if err == nil {
			// done writing list
			err = p.c.PrintfLine("")
		} else {
			err = p.Error(err.Error())
		}
		break
	case "LIST":
		if len(parts) == 2 {
			// 1 message
			idx, err = strconv.Atoi(parts[1])
			if err == nil {
				if idx > 0 && idx <= len(p.msgs) {
					info, err = os.Stat(p.msgs[idx-1].Filepath())
					if err == nil {
						err = p.c.PrintfLine("+OK %d %d", idx, info.Size())
					}
				} else {
					// no existing message
					err = p.Error("no such message")
				}
			}
		} else {
			// all messages
			err = p.c.PrintfLine("+OK %d messages (%d octets)", len(p.msgs), p.octs)
			if err == nil {
				dw := p.c.DotWriter()
				for i, msg := range p.msgs {
					info, err = os.Stat(msg.Filepath())
					if err == nil {
						// write out entry
						fmt.Fprintf(dw, "%d %d\r\n", i+1, info.Size())
					} else {
						log.Errorf("error in pop3, stat(): %s", err.Error())
					}
				}
				// FLUSH IT !!! :-DDDDD
				err = dw.Close()
			}
		}
		break
	default:
		err = p.Error("bad command")
	}
	return
}

// handle 1 line of input when not in transaction mode
func (p *pop3Session) handleLine(line string) (err error) {
	parts := strings.Split(line, " ")
	cmd := strings.ToUpper(parts[0])
	c := p.c
	switch cmd {
	case "STLS":
		if p.s.TLS == nil {
			err = p.Error("no TLS supported")
		} else {
			p.OK("Begin TLS negotiation")
			p.c, _, err = starttls.HandleStartTLS(p.nc, p.s.TLS)
			if err != nil {
				c.Close()
			}
		}
	case "NOOP":
		err = p.OK("")
		break
	case "PASS":
		if p.s.checkUser(p.user, line[5:]) {
			p.msgs, p.octs, err = p.s.obtainMessages(p.user)
			if err == nil {
				err = p.c.PrintfLine("+OK %s maildrop logged in, you have %d messages (%d octets)", p.user, len(p.msgs), p.octs)
				p.transaction = err == nil
			} else {
				err = p.Error(err.Error())
			}
		} else {
			err = p.Error("bad login")
		}
		break
	case "USER":
		if len(parts) > 1 {
			p.user = line[5:]
		}
		err = p.OK("you may try login as " + p.user)
		break
	default:
		err = p.Error("bad command")
	}
	return
}

func (p *pop3Session) OK(msg string) (err error) {
	if len(msg) == 0 {
		err = p.c.PrintfLine("+OK")
	} else {
		err = p.c.PrintfLine("+OK %s", msg)
	}
	return
}

// send error
func (p *pop3Session) Error(msg string) (err error) {
	err = p.c.PrintfLine("-ERR %s", msg)
	return
}

// serve sessions with connections accepted from a net.Listener
func (s *Server) Serve(l net.Listener) (err error) {
	for err == nil {
		var c net.Conn
		c, err = l.Accept()
		if err == nil {
			p := &pop3Session{
				nc: c,
				c:  textproto.NewConn(c),
				s:  s,
			}
			go p.Run()
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/fzzy/radix/redis"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"gopkg.in/vmihailenco/msgpack.v2"
	pb "proto"
	"testing"
)

const (
	address  = "localhost:50004"
	test_key = "testing:3721"
)

type TestStruct struct {
	Id   int32
	Name string
	Sex  int
	Data []byte
}

func TestBgSave(t *testing.T) {
	// start connection to redis
	client, err := redis.Dial("tcp", DEFAULT_REDIS_HOST)
	if err != nil {
		t.Fatal(err)
	}

	// mset data into redis
	bin, _ := msgpack.Marshal(&TestStruct{3721, "hello", 18, []byte("world")})
	reply := client.Cmd("set", test_key, bin)
	if reply.Err != nil {
		t.Fatal(reply.Err)
	}

	// Set up a connection to the server.
	conn, err := grpc.Dial(address)
	if err != nil {
		t.Fatal("did not connect: %v", err)
	}
	defer conn.Close()
	c := pb.NewBgSaveServiceClient(conn)

	// Contact the server and print out its response.
	_, err = c.MarkDirty(context.Background(), &pb.BgSave_Key{Name: test_key})
	if err != nil {
		t.Fatalf("could not query: %v", err)
	}
}
<commit_msg>add more test<commit_after>package main

import (
	log "github.com/GameGophers/nsq-logger"
	"github.com/fzzy/radix/redis"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/vmihailenco/msgpack.v2"
	pb "proto"
	"testing"
)

const (
	address  = "localhost:50004"
	test_key = "testing:3721"
)

type TestStruct struct {
	Id    int32
	Name  string
	Sex   int
	Int64 int64
	F64   float64
	F32   float32
	Data  []byte
}

func TestBgSave(t *testing.T) {
	return
	// start connection to redis
	client, err := redis.Dial("tcp", DEFAULT_REDIS_HOST)
	if err != nil {
		t.Fatal(err)
	}

	// mset data into redis
	bin, _ := msgpack.Marshal(&TestStruct{3721, "hello", 18, 999, 1.1, 2.2, []byte("world")})
	reply := client.Cmd("set", test_key, bin)
	if reply.Err != nil {
		t.Fatal(reply.Err)
	}

	// Set up a connection to the server.
	conn, err := grpc.Dial(address)
	if err != nil {
		t.Fatal("did not connect: %v", err)
	}
	defer conn.Close()
	c := pb.NewBgSaveServiceClient(conn)

	// Contact the server and print out its response.
	_, err = c.MarkDirty(context.Background(), &pb.BgSave_Key{Name: test_key})
	if err != nil {
		t.Fatalf("could not query: %v", err)
	}
}
func TestWriteBgSave(t *testing.T) {
	// start connection to mongodb
	sess, err := mgo.Dial(_mongodb_url)
	if err != nil {
		log.Critical(SERVICE, err)
		return
	}
	defer sess.Close()
	// database is provided in url
	db := sess.DB("test")
	tt := &TestStruct{3721, "hello", 18, 999, 1.1, 2.2, []byte("world")}
	_, err = db.C("testing").Upsert(bson.M{"Id": tt.Id}, tt)
	if err != nil {
		log.Critical(SERVICE, err)
		return
	}
	// get data from mongodb
	log.Info(SERVICE, "Save to mongodb: ", tt)
}
func TestReadBgSave(t *testing.T) {
	// start connection to mongodb
	sess, err := mgo.Dial(_mongodb_url)
	if err != nil {
		log.Critical(SERVICE, err)
		return
	}
	defer sess.Close()
	// database is provided in url
	db := sess.DB("test")
	tt := &TestStruct{}
	err = db.C("testing").Find(bson.M{"Id": int32(3721)}).One(&tt)
	if err != nil {
		log.Critical(SERVICE, err)
		return
	}
	// get data from mongodb
	log.Info(SERVICE, "Read from mongodb: ", tt)
}
<|endoftext|>