<commit_before>/*
Copyright 2015 Gravitational, Inc.

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 web

import (
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/gravitational/teleport/lib/auth"
	"github.com/gravitational/teleport/lib/services"
	"github.com/gravitational/teleport/lib/session"
	"github.com/gravitational/teleport/lib/utils"

	log "github.com/Sirupsen/logrus"
	"github.com/gravitational/trace"
	"github.com/gravitational/ttlmap"
)

// SessionContext is a context associated with users'
// web session, it stores connected client that persists
// between requests for example to avoid connecting
// to the auth server on every page hit
type SessionContext struct {
	sync.Mutex
	*log.Entry
	sess    *auth.Session
	user    string
	clt     *auth.TunClient
	parent  *sessionCache
	closers []io.Closer
}

func (c *SessionContext) getConnectHandler(sessionID session.ID) (*connectHandler, error) {
	c.Lock()
	defer c.Unlock()

	for _, closer := range c.closers {
		handler, ok := closer.(*connectHandler)
		if ok && handler.req.SessionID == sessionID {
			return handler, nil
		}
	}
	return nil, trace.NotFound("no connected streams")
}

func (c *SessionContext) UpdateSessionTerminal(sessionID session.ID, params session.TerminalParams) error {
	err := c.clt.UpdateSession(session.UpdateRequest{ID: sessionID, TerminalParams: &params})
	if err != nil {
		return trace.Wrap(err)
	}
	handler, err := c.getConnectHandler(sessionID)
	if err != nil {
		return trace.Wrap(err)
	}
	return trace.Wrap(handler.resizePTYWindow(params))
}

func (c *SessionContext) AddClosers(closers ...io.Closer) {
	c.Lock()
	defer c.Unlock()
	c.closers = append(c.closers, closers...)
}

func (c *SessionContext) TransferClosers() []io.Closer {
	c.Lock()
	defer c.Unlock()
	closers := c.closers
	c.closers = nil
	return closers
}

func (c *SessionContext) Invalidate() error {
	return c.parent.InvalidateSession(c)
}

// GetClient returns the client connected to the auth server
func (c *SessionContext) GetClient() (auth.ClientI, error) {
	return c.clt, nil
}

// GetUser returns the authenticated teleport user
func (c *SessionContext) GetUser() string {
	return c.user
}

// GetWebSession returns a web session
func (c *SessionContext) GetWebSession() *auth.Session {
	return c.sess
}

// ExtendWebSession creates a new web session for this user
// based on the previous session
func (c *SessionContext) ExtendWebSession() (*auth.Session, error) {
	sess, err := c.clt.ExtendWebSession(c.user, c.sess.ID)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return sess, nil
}

// GetAgent returns agent that can we used to answer challenges
// for the web to ssh connection
func (c *SessionContext) GetAgent() (auth.AgentCloser, error) {
	a, err := c.clt.GetAgent()
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return a, nil
}

// Close cleans up connections associated with requests
func (c *SessionContext) Close() error {
	closers := c.TransferClosers()
	for _, closer := range closers {
		c.Infof("closing %v", closer)
		closer.Close()
	}
	if c.clt != nil {
		return trace.Wrap(c.clt.Close())
	}
	return nil
}

// newSessionCache returns new instance of the session cache
func newSessionCache(servers []utils.NetAddr) (*sessionCache, error) {
	m, err := ttlmap.New(1024, ttlmap.CallOnExpire(closeContext))
	if err != nil {
		return nil, err
	}
	cache := &sessionCache{
		contexts:    m,
		authServers: servers,
		closer:      utils.NewCloseBroadcaster(),
	}
	// periodically close expired and unused sessions
	go cache.expireSessions()
	return cache, nil
}

// sessionCache handles web session authentication,
// and holds in memory contexts associated with each session
type sessionCache struct {
	sync.Mutex
	contexts    *ttlmap.TTLMap
	authServers []utils.NetAddr
	closer      *utils.CloseBroadcaster
}

// Close closes all allocated resources and stops goroutines
func (s *sessionCache) Close() error {
	log.Infof("[WEB] closing session cache")
	return s.closer.Close()
}

// closeContext is called when session context expires from
// cache and will clean up connections
func closeContext(key string, val interface{}) {
	go func() {
		log.Infof("[WEB] closing context %v", key)
		ctx, ok := val.(*SessionContext)
		if !ok {
			log.Warningf("warning, not valid value type %T", val)
			return
		}
		if err := ctx.Close(); err != nil {
			log.Infof("failed to close context: %v", err)
		}
	}()
}

func (s *sessionCache) expireSessions() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	go func() {
		for {
			select {
			case <-ticker.C:
				s.clearExpiredSessions()
			case <-s.closer.C:
				return
			}
		}
	}()
}

func (s *sessionCache) clearExpiredSessions() {
	s.Lock()
	defer s.Unlock()
	expired := s.contexts.RemoveExpired(10)
	if expired != 0 {
		log.Infof("[WEB] removed %v expired sessions", expired)
	}
}

func (s *sessionCache) Auth(user, pass string, hotpToken string) (*auth.Session, error) {
	method, err := auth.NewWebPasswordAuth(user, []byte(pass), hotpToken)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	clt, err := auth.NewTunClient("web.client", s.authServers, user, method)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	// we are always closing this client, because we will not be using
	// this connection initiated using password based credentials
	// down the road, so it's a one call client
	defer clt.Close()
	session, err := clt.SignIn(user, []byte(pass))
	if err != nil {
		defer clt.Close()
		return nil, trace.Wrap(err)
	}
	return session, nil
}

func (s *sessionCache) GetCertificate(c createSSHCertReq) (*SSHLoginResponse, error) {
	method, err := auth.NewWebPasswordAuth(c.User, []byte(c.Password),
		c.HOTPToken)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	clt, err := auth.NewTunClient("web.session", s.authServers, c.User, method)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	defer clt.Close()
	cert, err := clt.GenerateUserCert(c.PubKey, c.User, c.TTL)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	hostSigners, err := clt.GetCertAuthorities(services.HostCA, false)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	signers := []services.CertAuthority{}
	for _, hs := range hostSigners {
		signers = append(signers, *hs)
	}

	return &SSHLoginResponse{
		Cert:        cert,
		HostSigners: signers,
	}, nil
}

func (s *sessionCache) GetUserInviteInfo(token string) (user string,
	QRImg []byte, hotpFirstValues []string, e error) {

	method, err := auth.NewSignupTokenAuth(token)
	if err != nil {
		return "", nil, nil, trace.Wrap(err)
	}
	clt, err := auth.NewTunClient("web.get-user-invite", s.authServers, "tokenAuth", method)
	if err != nil {
		return "", nil, nil, trace.Wrap(err)
	}
	defer clt.Close()

	return clt.GetSignupTokenData(token)
}

func (s *sessionCache) CreateNewUser(token, password, hotpToken string) (*auth.Session, error) {
	method, err := auth.NewSignupTokenAuth(token)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	clt, err := auth.NewTunClient("web.create-user", s.authServers, "tokenAuth", method)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	defer clt.Close()
	sess, err := clt.CreateUserWithToken(token, password, hotpToken)
	return sess, trace.Wrap(err)
}

func (s *sessionCache) InvalidateSession(ctx *SessionContext) error {
	defer ctx.Close()
	if err := s.resetContext(ctx.GetUser(), ctx.GetWebSession().ID); err != nil {
		return trace.Wrap(err)
	}
	clt, err := ctx.GetClient()
	if err != nil {
		return trace.Wrap(err)
	}
	err = clt.DeleteWebSession(ctx.GetUser(), ctx.GetWebSession().ID)
	return trace.Wrap(err)
}

func (s *sessionCache) getContext(user, sid string) (*SessionContext, error) {
	s.Lock()
	defer s.Unlock()

	val, ok := s.contexts.Get(user + sid)
	if ok {
		return val.(*SessionContext), nil
	}
	return nil, trace.NotFound("sessionContext not found")
}

func (s *sessionCache) insertContext(user, sid string, ctx *SessionContext, ttl time.Duration) (*SessionContext, error) {
	s.Lock()
	defer s.Unlock()

	val, ok := s.contexts.Get(user + sid)
	if ok && val != nil { // nil means that we've just invalidated the context now and set it to nil in the cache
		return val.(*SessionContext), trace.AlreadyExists("exists")
	}
	if err := s.contexts.Set(user+sid, ctx, ttl); err != nil {
		return nil, trace.Wrap(err)
	}
	return ctx, nil
}

func (s *sessionCache) resetContext(user, sid string) error {
	s.Lock()
	defer s.Unlock()
	context, ok := s.contexts.Remove(user + sid)
	if ok {
		closeContext(user+sid, context)
	}
	return nil
}

func (s *sessionCache) ValidateSession(user, sid string) (*SessionContext, error) {
	ctx, err := s.getContext(user, sid)
	if err == nil {
		return ctx, nil
	}
	log.Debugf("ValidateSession(%s, %s)", user, sid)
	method, err := auth.NewWebSessionAuth(user, []byte(sid))
	if err != nil {
		return nil, trace.Wrap(err)
	}
	// Note: do not close this auth API client now. It will exist inside of "session context"
	clt, err := auth.NewTunClient("web.session-user", s.authServers, user, method)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	sess, err := clt.GetWebSessionInfo(user, sid)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	c := &SessionContext{
		clt:    clt,
		user:   user,
		sess:   sess,
		parent: s,
	}
	c.Entry = log.WithFields(log.Fields{
		"user": user,
		"sess": sess.ID[:4],
	})

	out, err := s.insertContext(user, sid, c, auth.WebSessionTTL)
	if err != nil {
		// this means that someone has just inserted the context, so
		// close our extra context and return
		if trace.IsAlreadyExists(err) {
			log.Infof("just created, returning the existing one")
			defer c.Close()
			return out, nil
		}
		return nil, trace.Wrap(err)
	}
	return out, nil
}

func (s *sessionCache) SetSession(w http.ResponseWriter, user, sid string) error {
	d, err := EncodeCookie(user, sid)
	if err != nil {
		return err
	}
	c := &http.Cookie{
		Name:     "session",
		Value:    d,
		Path:     "/",
		HttpOnly: true,
		Secure:   true,
	}
	http.SetCookie(w, c)
	return nil
}

func (s *sessionCache) ClearSession(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     "session",
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		Secure:   true,
	})
}
<commit_msg>Fix expired session cleanup<commit_after>/*
Copyright 2015 Gravitational, Inc.

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 web

import (
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/gravitational/teleport/lib/auth"
	"github.com/gravitational/teleport/lib/services"
	"github.com/gravitational/teleport/lib/session"
	"github.com/gravitational/teleport/lib/utils"

	log "github.com/Sirupsen/logrus"
	"github.com/gravitational/trace"
	"github.com/gravitational/ttlmap"
)

// SessionContext is a context associated with users'
// web session, it stores connected client that persists
// between requests for example to avoid connecting
// to the auth server on every page hit
type SessionContext struct {
	sync.Mutex
	*log.Entry
	sess    *auth.Session
	user    string
	clt     *auth.TunClient
	parent  *sessionCache
	closers []io.Closer
}

func (c *SessionContext) getConnectHandler(sessionID session.ID) (*connectHandler, error) {
	c.Lock()
	defer c.Unlock()

	for _, closer := range c.closers {
		handler, ok := closer.(*connectHandler)
		if ok && handler.req.SessionID == sessionID {
			return handler, nil
		}
	}
	return nil, trace.NotFound("no connected streams")
}

func (c *SessionContext) UpdateSessionTerminal(sessionID session.ID, params session.TerminalParams) error {
	err := c.clt.UpdateSession(session.UpdateRequest{ID: sessionID, TerminalParams: &params})
	if err != nil {
		return trace.Wrap(err)
	}
	handler, err := c.getConnectHandler(sessionID)
	if err != nil {
		return trace.Wrap(err)
	}
	return trace.Wrap(handler.resizePTYWindow(params))
}

func (c *SessionContext) AddClosers(closers ...io.Closer) {
	c.Lock()
	defer c.Unlock()
	c.closers = append(c.closers, closers...)
}

func (c *SessionContext) TransferClosers() []io.Closer {
	c.Lock()
	defer c.Unlock()
	closers := c.closers
	c.closers = nil
	return closers
}

func (c *SessionContext) Invalidate() error {
	return c.parent.InvalidateSession(c)
}

// GetClient returns the client connected to the auth server
func (c *SessionContext) GetClient() (auth.ClientI, error) {
	return c.clt, nil
}

// GetUser returns the authenticated teleport user
func (c *SessionContext) GetUser() string {
	return c.user
}

// GetWebSession returns a web session
func (c *SessionContext) GetWebSession() *auth.Session {
	return c.sess
}

// ExtendWebSession creates a new web session for this user
// based on the previous session
func (c *SessionContext) ExtendWebSession() (*auth.Session, error) {
	sess, err := c.clt.ExtendWebSession(c.user, c.sess.ID)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return sess, nil
}

// GetAgent returns agent that can we used to answer challenges
// for the web to ssh connection
func (c *SessionContext) GetAgent() (auth.AgentCloser, error) {
	a, err := c.clt.GetAgent()
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return a, nil
}

// Close cleans up connections associated with requests
func (c *SessionContext) Close() error {
	closers := c.TransferClosers()
	for _, closer := range closers {
		c.Infof("closing %v", closer)
		closer.Close()
	}
	if c.clt != nil {
		return trace.Wrap(c.clt.Close())
	}
	return nil
}

// newSessionCache returns new instance of the session cache
func newSessionCache(servers []utils.NetAddr) (*sessionCache, error) {
	m, err := ttlmap.New(1024, ttlmap.CallOnExpire(closeContext))
	if err != nil {
		return nil, err
	}
	cache := &sessionCache{
		contexts:    m,
		authServers: servers,
		closer:      utils.NewCloseBroadcaster(),
	}
	// periodically close expired and unused sessions
	go cache.expireSessions()
	return cache, nil
}

// sessionCache handles web session authentication,
// and holds in memory contexts associated with each session
type sessionCache struct {
	sync.Mutex
	contexts    *ttlmap.TTLMap
	authServers []utils.NetAddr
	closer      *utils.CloseBroadcaster
}

// Close closes all allocated resources and stops goroutines
func (s *sessionCache) Close() error {
	log.Infof("[WEB] closing session cache")
	return s.closer.Close()
}

// closeContext is called when session context expires from
// cache and will clean up connections
func closeContext(key string, val interface{}) {
	go func() {
		log.Infof("[WEB] closing context %v", key)
		ctx, ok := val.(*SessionContext)
		if !ok {
			log.Warningf("warning, not valid value type %T", val)
			return
		}
		if err := ctx.Close(); err != nil {
			log.Infof("failed to close context: %v", err)
		}
	}()
}

func (s *sessionCache) expireSessions() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			s.clearExpiredSessions()
		case <-s.closer.C:
			return
		}
	}
}

func (s *sessionCache) clearExpiredSessions() {
	s.Lock()
	defer s.Unlock()
	expired := s.contexts.RemoveExpired(10)
	if expired != 0 {
		log.Infof("[WEB] removed %v expired sessions", expired)
	}
}

func (s *sessionCache) Auth(user, pass string, hotpToken string) (*auth.Session, error) {
	method, err := auth.NewWebPasswordAuth(user, []byte(pass), hotpToken)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	clt, err := auth.NewTunClient("web.client", s.authServers, user, method)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	// we are always closing this client, because we will not be using
	// this connection initiated using password based credentials
	// down the road, so it's a one call client
	defer clt.Close()
	session, err := clt.SignIn(user, []byte(pass))
	if err != nil {
		defer clt.Close()
		return nil, trace.Wrap(err)
	}
	return session, nil
}

func (s *sessionCache) GetCertificate(c createSSHCertReq) (*SSHLoginResponse, error) {
	method, err := auth.NewWebPasswordAuth(c.User, []byte(c.Password),
		c.HOTPToken)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	clt, err := auth.NewTunClient("web.session", s.authServers, c.User, method)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	defer clt.Close()
	cert, err := clt.GenerateUserCert(c.PubKey, c.User, c.TTL)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	hostSigners, err := clt.GetCertAuthorities(services.HostCA, false)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	signers := []services.CertAuthority{}
	for _, hs := range hostSigners {
		signers = append(signers, *hs)
	}

	return &SSHLoginResponse{
		Cert:        cert,
		HostSigners: signers,
	}, nil
}

func (s *sessionCache) GetUserInviteInfo(token string) (user string,
	QRImg []byte, hotpFirstValues []string, e error) {

	method, err := auth.NewSignupTokenAuth(token)
	if err != nil {
		return "", nil, nil, trace.Wrap(err)
	}
	clt, err := auth.NewTunClient("web.get-user-invite", s.authServers, "tokenAuth", method)
	if err != nil {
		return "", nil, nil, trace.Wrap(err)
	}
	defer clt.Close()

	return clt.GetSignupTokenData(token)
}

func (s *sessionCache) CreateNewUser(token, password, hotpToken string) (*auth.Session, error) {
	method, err := auth.NewSignupTokenAuth(token)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	clt, err := auth.NewTunClient("web.create-user", s.authServers, "tokenAuth", method)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	defer clt.Close()
	sess, err := clt.CreateUserWithToken(token, password, hotpToken)
	return sess, trace.Wrap(err)
}

func (s *sessionCache) InvalidateSession(ctx *SessionContext) error {
	defer ctx.Close()
	if err := s.resetContext(ctx.GetUser(), ctx.GetWebSession().ID); err != nil {
		return trace.Wrap(err)
	}
	clt, err := ctx.GetClient()
	if err != nil {
		return trace.Wrap(err)
	}
	err = clt.DeleteWebSession(ctx.GetUser(), ctx.GetWebSession().ID)
	return trace.Wrap(err)
}

func (s *sessionCache) getContext(user, sid string) (*SessionContext, error) {
	s.Lock()
	defer s.Unlock()

	val, ok := s.contexts.Get(user + sid)
	if ok {
		return val.(*SessionContext), nil
	}
	return nil, trace.NotFound("sessionContext not found")
}

func (s *sessionCache) insertContext(user, sid string, ctx *SessionContext, ttl time.Duration) (*SessionContext, error) {
	s.Lock()
	defer s.Unlock()

	val, ok := s.contexts.Get(user + sid)
	if ok && val != nil { // nil means that we've just invalidated the context now and set it to nil in the cache
		return val.(*SessionContext), trace.AlreadyExists("exists")
	}
	if err := s.contexts.Set(user+sid, ctx, ttl); err != nil {
		return nil, trace.Wrap(err)
	}
	return ctx, nil
}

func (s *sessionCache) resetContext(user, sid string) error {
	s.Lock()
	defer s.Unlock()
	context, ok := s.contexts.Remove(user + sid)
	if ok {
		closeContext(user+sid, context)
	}
	return nil
}

func (s *sessionCache) ValidateSession(user, sid string) (*SessionContext, error) {
	ctx, err := s.getContext(user, sid)
	if err == nil {
		return ctx, nil
	}
	log.Debugf("ValidateSession(%s, %s)", user, sid)
	method, err := auth.NewWebSessionAuth(user, []byte(sid))
	if err != nil {
		return nil, trace.Wrap(err)
	}
	// Note: do not close this auth API client now. It will exist inside of "session context"
	clt, err := auth.NewTunClient("web.session-user", s.authServers, user, method)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	sess, err := clt.GetWebSessionInfo(user, sid)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	c := &SessionContext{
		clt:    clt,
		user:   user,
		sess:   sess,
		parent: s,
	}
	c.Entry = log.WithFields(log.Fields{
		"user": user,
		"sess": sess.ID[:4],
	})

	out, err := s.insertContext(user, sid, c, auth.WebSessionTTL)
	if err != nil {
		// this means that someone has just inserted the context, so
		// close our extra context and return
		if trace.IsAlreadyExists(err) {
			log.Infof("just created, returning the existing one")
			defer c.Close()
			return out, nil
		}
		return nil, trace.Wrap(err)
	}
	return out, nil
}

func (s *sessionCache) SetSession(w http.ResponseWriter, user, sid string) error {
	d, err := EncodeCookie(user, sid)
	if err != nil {
		return err
	}
	c := &http.Cookie{
		Name:     "session",
		Value:    d,
		Path:     "/",
		HttpOnly: true,
		Secure:   true,
	}
	http.SetCookie(w, c)
	return nil
}

func (s *sessionCache) ClearSession(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     "session",
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		Secure:   true,
	})
}
<|endoftext|>


<commit_before>package logical

func SanitizeResponse(input *Response) *HTTPResponse {
	logicalResp := &HTTPResponse{
		Data:     input.Data,
		Warnings: input.Warnings(),
	}

	if input.Secret != nil {
		logicalResp.LeaseID = input.Secret.LeaseID
		logicalResp.Renewable = input.Secret.Renewable
		logicalResp.LeaseDuration = int(input.Secret.TTL.Seconds())
	}

	// If we have authentication information, then
	// set up the result structure.
	if input.Auth != nil {
		logicalResp.Auth = &HTTPAuth{
			ClientToken:   input.Auth.ClientToken,
			Accessor:      input.Auth.Accessor,
			Policies:      input.Auth.Policies,
			Metadata:      input.Auth.Metadata,
			LeaseDuration: int(input.Auth.TTL.Seconds()),
			Renewable:     input.Auth.Renewable,
		}
	}

	return logicalResp
}

type HTTPResponse struct {
	LeaseID       string                 `json:"lease_id"`
	Renewable     bool                   `json:"renewable"`
	LeaseDuration int                    `json:"lease_duration"`
	Data          map[string]interface{} `json:"data"`
	WrapInfo      *HTTPWrapInfo          `json:"wrap_info"`
	Warnings      []string               `json:"warnings"`
	Auth          *HTTPAuth              `json:"auth"`
}

type HTTPAuth struct {
	ClientToken   string            `json:"client_token"`
	Accessor      string            `json:"accessor"`
	Policies      []string          `json:"policies"`
	Metadata      map[string]string `json:"metadata"`
	LeaseDuration int               `json:"lease_duration"`
	Renewable     bool              `json:"renewable"`
}

type HTTPWrapInfo struct {
	Token string `json:"token"`
	TTL   int    `json:"ttl"`
}
<commit_msg>Add some comments to sanitize<commit_after>package logical

// This logic was pulled from the http package so that it can be used for
// encoding wrapped responses as well. It simply translates the logical request
// to an http response, with the values we want and omitting the values we
// don't.
func SanitizeResponse(input *Response) *HTTPResponse {
	logicalResp := &HTTPResponse{
		Data:     input.Data,
		Warnings: input.Warnings(),
	}

	if input.Secret != nil {
		logicalResp.LeaseID = input.Secret.LeaseID
		logicalResp.Renewable = input.Secret.Renewable
		logicalResp.LeaseDuration = int(input.Secret.TTL.Seconds())
	}

	// If we have authentication information, then
	// set up the result structure.
	if input.Auth != nil {
		logicalResp.Auth = &HTTPAuth{
			ClientToken:   input.Auth.ClientToken,
			Accessor:      input.Auth.Accessor,
			Policies:      input.Auth.Policies,
			Metadata:      input.Auth.Metadata,
			LeaseDuration: int(input.Auth.TTL.Seconds()),
			Renewable:     input.Auth.Renewable,
		}
	}

	return logicalResp
}

type HTTPResponse struct {
	LeaseID       string                 `json:"lease_id"`
	Renewable     bool                   `json:"renewable"`
	LeaseDuration int                    `json:"lease_duration"`
	Data          map[string]interface{} `json:"data"`
	WrapInfo      *HTTPWrapInfo          `json:"wrap_info"`
	Warnings      []string               `json:"warnings"`
	Auth          *HTTPAuth              `json:"auth"`
}

type HTTPAuth struct {
	ClientToken   string            `json:"client_token"`
	Accessor      string            `json:"accessor"`
	Policies      []string          `json:"policies"`
	Metadata      map[string]string `json:"metadata"`
	LeaseDuration int               `json:"lease_duration"`
	Renewable     bool              `json:"renewable"`
}

type HTTPWrapInfo struct {
	Token string `json:"token"`
	TTL   int    `json:"ttl"`
}
<|endoftext|>


<commit_before>package lzma

import "io"

type bWriter struct {
	io.Writer
	a []byte
}

func newByteWriter(w io.Writer) io.ByteWriter {
	if b, ok := w.(io.ByteWriter); ok {
		return b
	}
	return &bWriter{w, make([]byte, 1)}
}

func (b *bWriter) WriteByte(c byte) error {
	b.a[0] = c
	n, err := b.Write(b.a)
	switch {
	case n > 1:
		panic("n > 1 for writing a single byte")
	case n == 1:
		return nil
	case err == nil:
		panic("no error for n == 0")
	}
	return err
}

type bReader struct {
	io.Reader
	a []byte
}

func newByteReader(r io.Reader) io.ByteReader {
	if b, ok := r.(io.ByteReader); ok {
		return b
	}
	return &bReader{r, make([]byte, 1)}
}

func (b bReader) ReadByte() (byte, error) {
	n, err := b.Read(b.a)
	switch {
	case n > 1:
		panic("n < 1 for reading a single byte")
	case n == 1:
		return b.a[0], nil
	}
	return 0, err
}

// rangeEncoder implements range encoding of single bits. The low value can
// overflow therefore we need uint64. The cache value is used to handle
// overflows.
type rangeEncoder struct {
	w         io.ByteWriter
	range_    uint32
	low       uint64
	cacheSize int64
	cache     byte
	// for debugging
	// bitCounter int
}

// newRangeEncoder creates a new range encoder.
func newRangeEncoder(w io.Writer) *rangeEncoder {
	return &rangeEncoder{
		w:         newByteWriter(w),
		range_:    0xffffffff,
		cacheSize: 1}
}

// DirectEncodeBit encodes the least-significant bit of b with probability 1/2.
func (e *rangeEncoder) DirectEncodeBit(b uint32) error {
	// e.bitCounter++
	e.range_ >>= 1
	e.low += uint64(e.range_) & (0 - (uint64(b) & 1))
	if err := e.normalize(); err != nil {
		return err
	}

	// debug.Printf("D %3d %0x08x %d\n", e.bitCounter, e.range_, b)
	return nil
}

// EncodeBit encodes the least significant bit of b. The p value will be
// updated by the function depending on the bit encoded.
func (e *rangeEncoder) EncodeBit(b uint32, p *prob) error {
	// e.bitCounter++
	bound := p.bound(e.range_)
	if b&1 == 0 {
		e.range_ = bound
		p.inc()
	} else {
		e.low += uint64(bound)
		e.range_ -= bound
		p.dec()
	}
	if err := e.normalize(); err != nil {
		return err
	}

	// debug.Printf("B %3d 0x%08x 0x%03x %d\n", e.bitCounter, e.range_, *p, b)
	return nil
}

// Flush writes a complete copy of the low value.
func (e *rangeEncoder) Close() error {
	for i := 0; i < 5; i++ {
		if err := e.shiftLow(); err != nil {
			return err
		}
	}
	return nil
}

// newRangeDecoder initializes a range decoder. It reads five bytes from the
// reader and therefore may return an error.
func newRangeDecoder(r io.Reader) (d *rangeDecoder, err error) {
	d = &rangeDecoder{r: newByteReader(r)}
	err = d.init()
	return
}

// possiblyAtEnd checks whether the decoder may be at the end of the stream.
func (d *rangeDecoder) possiblyAtEnd() bool {
	return d.code == 0
}

var bitCounter int

// DirectDecodeBit decodes a bit with probability 1/2. The return value b will
// contain the bit at the least-significant position. All other bits will be
// zero.
func (d *rangeDecoder) DirectDecodeBit() (b uint32, err error) {
	// d.bitCounter++
	d.range_ >>= 1
	d.code -= d.range_
	t := 0 - (d.code >> 31)
	d.code += d.range_ & t

	// d.code will stay less then d.range_

	if err = d.normalize(); err != nil {
		return 0, err
	}

	b = (t + 1) & 1

	// debug.Printf("D %3d 0x%08x %d\n", d.bitCounter, d.range_, b)
	return b, nil
}

// decodeBit decodes a single bit. The bit will be returned at the
// least-significant position. All other bits will be zero. The probability
// value will be updated.
func (d *rangeDecoder) DecodeBit(p *prob) (b uint32, err error) {
	// d.bitCounter++
	bound := p.bound(d.range_)
	if d.code < bound {
		d.range_ = bound
		p.inc()
		b = 0
	} else {
		d.code -= bound
		d.range_ -= bound
		p.dec()
		b = 1
	}

	// d.code will stay less then d.range_

	if err = d.normalize(); err != nil {
		return 0, err
	}

	// debug.Printf("B %3d 0x%08x 0x%03x %d\n", d.bitCounter, d.range_, *p, b)
	return b, nil
}

// shiftLow() shifts the low value for 8 bit. The shifted byte is written into
// the byte writer. The cache value is used to handle overflows.
func (e *rangeEncoder) shiftLow() error {
	if uint32(e.low) < 0xff000000 || (e.low>>32) != 0 {
		tmp := e.cache
		for {
			err := e.w.WriteByte(tmp + byte(e.low>>32))
			if err != nil {
				return err
			}
			tmp = 0xff
			e.cacheSize--
			if e.cacheSize <= 0 {
				if e.cacheSize < 0 {
					return newError("negative e.cacheSize")
				}
				break
			}
		}
		e.cache = byte(uint32(e.low) >> 24)
	}
	e.cacheSize++
	e.low = uint64(uint32(e.low) << 8)
	return nil
}

// normalize handles shifts of range_ and low.
func (e *rangeEncoder) normalize() error {
	const top = 1 << 24
	if e.range_ >= top {
		return nil
	}
	e.range_ <<= 8
	return e.shiftLow()
}

// rangeDecoder decodes single bits of the range encoding stream.
type rangeDecoder struct {
	r      io.ByteReader
	range_ uint32
	code   uint32
	// for debugging
	// bitCounter int
}

// init initializes the range decoder, by reading from the byte reader.
func (d *rangeDecoder) init() error {
	d.range_ = 0xffffffff
	d.code = 0

	b, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	if b != 0 {
		return newError("first byte not zero")
	}

	for i := 0; i < 4; i++ {
		if err = d.updateCode(); err != nil {
			return err
		}
	}

	if d.code >= d.range_ {
		return newError("newRangeDecoder: d.code >= d.range_")
	}

	return nil
}

// updateCode reads a new byte into the code.
func (d *rangeDecoder) updateCode() error {
	b, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	d.code = (d.code << 8) | uint32(b)
	return nil
}

// normalize the top value and update the code value.
func (d *rangeDecoder) normalize() error {
	// assume d.code < d.range_
	const top = 1 << 24
	if d.range_ < top {
		d.range_ <<= 8
		// d.code < d.range_ will be maintained
		if err := d.updateCode(); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>lzma: removed bitCounter from range_codec.go<commit_after>package lzma

import "io"

type bWriter struct {
	io.Writer
	a []byte
}

func newByteWriter(w io.Writer) io.ByteWriter {
	if b, ok := w.(io.ByteWriter); ok {
		return b
	}
	return &bWriter{w, make([]byte, 1)}
}

func (b *bWriter) WriteByte(c byte) error {
	b.a[0] = c
	n, err := b.Write(b.a)
	switch {
	case n > 1:
		panic("n > 1 for writing a single byte")
	case n == 1:
		return nil
	case err == nil:
		panic("no error for n == 0")
	}
	return err
}

type bReader struct {
	io.Reader
	a []byte
}

func newByteReader(r io.Reader) io.ByteReader {
	if b, ok := r.(io.ByteReader); ok {
		return b
	}
	return &bReader{r, make([]byte, 1)}
}

func (b bReader) ReadByte() (byte, error) {
	n, err := b.Read(b.a)
	switch {
	case n > 1:
		panic("n < 1 for reading a single byte")
	case n == 1:
		return b.a[0], nil
	}
	return 0, err
}

// rangeEncoder implements range encoding of single bits. The low value can
// overflow therefore we need uint64. The cache value is used to handle
// overflows.
type rangeEncoder struct {
	w         io.ByteWriter
	range_    uint32
	low       uint64
	cacheSize int64
	cache     byte
	// for debugging
	// bitCounter int
}

// newRangeEncoder creates a new range encoder.
func newRangeEncoder(w io.Writer) *rangeEncoder {
	return &rangeEncoder{
		w:         newByteWriter(w),
		range_:    0xffffffff,
		cacheSize: 1}
}

// DirectEncodeBit encodes the least-significant bit of b with probability 1/2.
func (e *rangeEncoder) DirectEncodeBit(b uint32) error {
	// e.bitCounter++
	e.range_ >>= 1
	e.low += uint64(e.range_) & (0 - (uint64(b) & 1))
	if err := e.normalize(); err != nil {
		return err
	}

	// debug.Printf("D %3d %0x08x %d\n", e.bitCounter, e.range_, b)
	return nil
}

// EncodeBit encodes the least significant bit of b. The p value will be
// updated by the function depending on the bit encoded.
func (e *rangeEncoder) EncodeBit(b uint32, p *prob) error {
	// e.bitCounter++
	bound := p.bound(e.range_)
	if b&1 == 0 {
		e.range_ = bound
		p.inc()
	} else {
		e.low += uint64(bound)
		e.range_ -= bound
		p.dec()
	}
	if err := e.normalize(); err != nil {
		return err
	}

	// debug.Printf("B %3d 0x%08x 0x%03x %d\n", e.bitCounter, e.range_, *p, b)
	return nil
}

// Flush writes a complete copy of the low value.
func (e *rangeEncoder) Close() error {
	for i := 0; i < 5; i++ {
		if err := e.shiftLow(); err != nil {
			return err
		}
	}
	return nil
}

// newRangeDecoder initializes a range decoder. It reads five bytes from the
// reader and therefore may return an error.
func newRangeDecoder(r io.Reader) (d *rangeDecoder, err error) {
	d = &rangeDecoder{r: newByteReader(r)}
	err = d.init()
	return
}

// possiblyAtEnd checks whether the decoder may be at the end of the stream.
func (d *rangeDecoder) possiblyAtEnd() bool {
	return d.code == 0
}

// DirectDecodeBit decodes a bit with probability 1/2. The return value b will
// contain the bit at the least-significant position. All other bits will be
// zero.
func (d *rangeDecoder) DirectDecodeBit() (b uint32, err error) {
	// d.bitCounter++
	d.range_ >>= 1
	d.code -= d.range_
	t := 0 - (d.code >> 31)
	d.code += d.range_ & t

	// d.code will stay less then d.range_

	if err = d.normalize(); err != nil {
		return 0, err
	}

	b = (t + 1) & 1

	// debug.Printf("D %3d 0x%08x %d\n", d.bitCounter, d.range_, b)
	return b, nil
}

// decodeBit decodes a single bit. The bit will be returned at the
// least-significant position. All other bits will be zero. The probability
// value will be updated.
func (d *rangeDecoder) DecodeBit(p *prob) (b uint32, err error) {
	// d.bitCounter++
	bound := p.bound(d.range_)
	if d.code < bound {
		d.range_ = bound
		p.inc()
		b = 0
	} else {
		d.code -= bound
		d.range_ -= bound
		p.dec()
		b = 1
	}

	// d.code will stay less then d.range_

	if err = d.normalize(); err != nil {
		return 0, err
	}

	// debug.Printf("B %3d 0x%08x 0x%03x %d\n", d.bitCounter, d.range_, *p, b)
	return b, nil
}

// shiftLow() shifts the low value for 8 bit. The shifted byte is written into
// the byte writer. The cache value is used to handle overflows.
func (e *rangeEncoder) shiftLow() error {
	if uint32(e.low) < 0xff000000 || (e.low>>32) != 0 {
		tmp := e.cache
		for {
			err := e.w.WriteByte(tmp + byte(e.low>>32))
			if err != nil {
				return err
			}
			tmp = 0xff
			e.cacheSize--
			if e.cacheSize <= 0 {
				if e.cacheSize < 0 {
					return newError("negative e.cacheSize")
				}
				break
			}
		}
		e.cache = byte(uint32(e.low) >> 24)
	}
	e.cacheSize++
	e.low = uint64(uint32(e.low) << 8)
	return nil
}

// normalize handles shifts of range_ and low.
func (e *rangeEncoder) normalize() error {
	const top = 1 << 24
	if e.range_ >= top {
		return nil
	}
	e.range_ <<= 8
	return e.shiftLow()
}

// rangeDecoder decodes single bits of the range encoding stream.
type rangeDecoder struct {
	r      io.ByteReader
	range_ uint32
	code   uint32
	// for debugging
	// bitCounter int
}

// init initializes the range decoder, by reading from the byte reader.
func (d *rangeDecoder) init() error {
	d.range_ = 0xffffffff
	d.code = 0

	b, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	if b != 0 {
		return newError("first byte not zero")
	}

	for i := 0; i < 4; i++ {
		if err = d.updateCode(); err != nil {
			return err
		}
	}

	if d.code >= d.range_ {
		return newError("newRangeDecoder: d.code >= d.range_")
	}

	return nil
}

// updateCode reads a new byte into the code.
func (d *rangeDecoder) updateCode() error {
	b, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	d.code = (d.code << 8) | uint32(b)
	return nil
}

// normalize the top value and update the code value.
func (d *rangeDecoder) normalize() error {
	// assume d.code < d.range_
	const top = 1 << 24
	if d.range_ < top {
		d.range_ <<= 8
		// d.code < d.range_ will be maintained
		if err := d.updateCode(); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>//
// gomanta - Go library to interact with Joyent Manta
//
//
// Copyright (c) 2013 Joyent Inc.
//
// Written by Daniele Stroppa <daniele.stroppa@joyent.com>
//

package manta_test

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"

	gc "launchpad.net/gocheck"

	"github.com/joyent/gocommon/client"
	localmanta "github.com/joyent/gomanta/localservices/manta"
	"github.com/joyent/gomanta/manta"
	"github.com/joyent/gosign/auth"
)

var privateKey []byte

func registerLocalTests(keyName string) {
	var localKeyFile string
	if keyName == "" {
		localKeyFile = os.Getenv("HOME") + "/.ssh/id_rsa"
	} else {
		localKeyFile = keyName
	}
	privateKey, _ = ioutil.ReadFile(localKeyFile)

	gc.Suite(&LocalTests{})
}

type LocalTests struct {
	LiveTests
	Server     *httptest.Server
	Mux        *http.ServeMux
	oldHandler http.Handler
	manta      *localmanta.Manta
}

func (s *LocalTests) SetUpSuite(c *gc.C) {
	// Set up the HTTP server.
	s.Server = httptest.NewServer(nil)
	s.oldHandler = s.Server.Config.Handler
	s.Mux = http.NewServeMux()
	s.Server.Config.Handler = s.Mux

	// Set up a Joyent Manta service.
	authentication := auth.Auth{User: "localtest", PrivateKey: string(privateKey), Algorithm: "rsa-sha256"}

	s.creds = &auth.Credentials{
		UserAuthentication: authentication,
		MantaKeyId:         "",
		MantaEndpoint:      auth.Endpoint{URL: s.Server.URL},
	}
	s.manta = localmanta.New(s.creds.MantaEndpoint.URL, s.creds.UserAuthentication.User)
	s.manta.SetupHTTP(s.Mux)
}

func (s *LocalTests) TearDownSuite(c *gc.C) {
	s.Mux = nil
	s.Server.Config.Handler = s.oldHandler
	s.Server.Close()
}

func (s *LocalTests) SetUpTest(c *gc.C) {
	client := client.NewClient(s.creds.MantaEndpoint.URL, "", s.creds, &manta.Logger)
	c.Assert(client, gc.NotNil)
	s.testClient = manta.New(client)
	c.Assert(s.testClient, gc.NotNil)
}

// Helper method to create a test directory
func (s *LocalTests) createDirectory(c *gc.C, path string) {
	err := s.testClient.PutDirectory(path)
	c.Assert(err, gc.IsNil)
}

// Helper method to create a test object
func (s *LocalTests) createObject(c *gc.C, path, objName string) {
	err := s.testClient.PutObject(path, objName, []byte("Test Manta API"))
	c.Assert(err, gc.IsNil)
}

// Helper method to delete a test directory
func (s *LocalTests) deleteDirectory(c *gc.C, path string) {
	err := s.testClient.DeleteDirectory(path)
	c.Assert(err, gc.IsNil)
}

// Helper method to delete a test object
func (s *LocalTests) deleteObject(c *gc.C, path, objName string) {
	err := s.testClient.DeleteObject(path, objName)
	c.Assert(err, gc.IsNil)
}

// Helper method to create a test job
func (s *LocalTests) createJob(c *gc.C, jobName string) string {
	phases := []manta.Phase{
		{Type: "map", Exec: "wc", Init: ""},
		{Type: "reduce", Exec: "awk '{ l += $1; w += $2; c += $3 } END { print l, w, c }'", Init: ""},
	}
	jobUri, err := s.testClient.CreateJob(manta.CreateJobOpts{Name: jobName, Phases: phases})
	c.Assert(err, gc.IsNil)
	c.Assert(jobUri, gc.NotNil)
	return strings.Split(jobUri, "/")[3]
}

// Storage API
func (s *LocalTests) TestPutDirectory(c *gc.C) {
	s.createDirectory(c, "test")

	// cleanup
	s.deleteDirectory(c, "test")
}

func (s *LocalTests) TestListDirectory(c *gc.C) {
	s.createDirectory(c, "test")
	defer s.deleteDirectory(c, "test")
	s.createObject(c, "test", "obj")
	defer s.deleteObject(c, "test", "obj")

	opts := manta.ListDirectoryOpts{}
	dirs, err := s.testClient.ListDirectory("test", opts)
	c.Assert(err, gc.IsNil)
	c.Assert(dirs, gc.NotNil)
}

func (s *LocalTests) TestDeleteDirectory(c *gc.C) {
	s.createDirectory(c, "test")

	s.deleteDirectory(c, "test")
}

func (s *LocalTests) TestPutObject(c *gc.C) {
	s.createDirectory(c, "dir")
	defer s.deleteDirectory(c, "dir")
	s.createObject(c, "dir", "obj")
	defer s.deleteObject(c, "dir", "obj")
}

func (s *LocalTests) TestGetObject(c *gc.C) {
	s.createDirectory(c, "dir")
	defer s.deleteDirectory(c, "dir")
	s.createObject(c, "dir", "obj")
	defer s.deleteObject(c, "dir", "obj")

	obj, err := s.testClient.GetObject("dir", "obj")
	c.Assert(err, gc.IsNil)
	c.Assert(obj, gc.NotNil)
	c.Check(string(obj), gc.Equals, "Test Manta API")
}

func (s *LocalTests) TestDeleteObject(c *gc.C) {
	s.createDirectory(c, "dir")
	defer s.deleteDirectory(c, "dir")
	s.createObject(c, "dir", "obj")

	s.deleteObject(c, "dir", "obj")
}

func (s *LocalTests) TestPutSnapLink(c *gc.C) {
	s.createDirectory(c, "linkdir")
	defer s.deleteDirectory(c, "linkdir")
	s.createObject(c, "linkdir", "obj")
	defer s.deleteObject(c, "linkdir", "obj")

	location := fmt.Sprintf("/%s/%s", s.creds.UserAuthentication.User, "stor/linkdir/obj")
	err := s.testClient.PutSnapLink("linkdir", "objlnk", location)
	c.Assert(err, gc.IsNil)

	// cleanup
	s.deleteObject(c, "linkdir", "objlnk")
}

// Jobs API
func (s *LocalTests) TestCreateJob(c *gc.C) {
	s.createJob(c, "test-job")
}

func (s *LocalTests) TestListLiveJobs(c *gc.C) {
	s.createJob(c, "test-job")

	jobs, err := s.testClient.ListJobs(true)
	c.Assert(err, gc.IsNil)
	c.Assert(jobs, gc.NotNil)
	c.Assert(len(jobs) >= 1, gc.Equals, true)
}

func (s *LocalTests) TestListAllJobs(c *gc.C) {
	s.createJob(c, "test-job")

	jobs, err := s.testClient.ListJobs(false)
	c.Assert(err, gc.IsNil)
	c.Assert(jobs, gc.NotNil)
}

func (s *LocalTests) TestCancelJob(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	err := s.testClient.CancelJob(jobId)
	c.Assert(err, gc.IsNil)
}

func (s *LocalTests) TestAddJobInputs(c *gc.C) {
	var inputs = `/` + s.creds.UserAuthentication.User + `/stor/testjob/obj1
/` + s.creds.UserAuthentication.User + `/stor/testjob/obj2
`
	s.createDirectory(c, "testjob")
	defer s.deleteDirectory(c, "testjob")
	s.createObject(c, "testjob", "obj1")
	defer s.deleteObject(c, "testjob", "obj1")
	s.createObject(c, "testjob", "obj2")
	defer s.deleteObject(c, "testjob", "obj2")
	jobId := s.createJob(c, "test-job")

	err := s.testClient.AddJobInputs(jobId, strings.NewReader(inputs))
	c.Assert(err, gc.IsNil)
}

func (s *LocalTests) TestEndJobInputs(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	err := s.testClient.EndJobInputs(jobId)
	c.Assert(err, gc.IsNil)
}

func (s *LocalTests) TestGetJob(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	job, err := s.testClient.GetJob(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(job, gc.NotNil)
}

func (s *LocalTests) TestGetJobInput(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	input, err := s.testClient.GetJobInput(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(input, gc.NotNil)
}

func (s *LocalTests) TestGetJobOutput(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	output, err := s.testClient.GetJobOutput(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(output, gc.NotNil)
}

func (s *LocalTests) TestGetJobFailures(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	fail, err := s.testClient.GetJobFailures(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(fail, gc.Equals, "")
}

func (s *LocalTests) TestGetJobErrors(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	errs, err := s.testClient.GetJobErrors(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(errs, gc.IsNil)
}
<commit_msg>Cache parsed signing key for performance<commit_after>//
// gomanta - Go library to interact with Joyent Manta
//
//
// Copyright (c) 2013 Joyent Inc.
//
// Written by Daniele Stroppa <daniele.stroppa@joyent.com>
//

package manta_test

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"

	gc "launchpad.net/gocheck"

	"github.com/joyent/gocommon/client"
	localmanta "github.com/joyent/gomanta/localservices/manta"
	"github.com/joyent/gomanta/manta"
	"github.com/joyent/gosign/auth"
)

var privateKey []byte

func registerLocalTests(keyName string) {
	var localKeyFile string
	if keyName == "" {
		localKeyFile = os.Getenv("HOME") + "/.ssh/id_rsa"
	} else {
		localKeyFile = keyName
	}
	privateKey, _ = ioutil.ReadFile(localKeyFile)

	gc.Suite(&LocalTests{})
}

type LocalTests struct {
	LiveTests
	Server     *httptest.Server
	Mux        *http.ServeMux
	oldHandler http.Handler
	manta      *localmanta.Manta
}

func (s *LocalTests) SetUpSuite(c *gc.C) {
	// Set up the HTTP server.
	s.Server = httptest.NewServer(nil)
	s.oldHandler = s.Server.Config.Handler
	s.Mux = http.NewServeMux()
	s.Server.Config.Handler = s.Mux

	// Set up a Joyent Manta service.
	authentication, err := auth.NewAuth("localtest", string(privateKey), "rsa-sha256")
	c.Assert(err, gc.IsNil)

	s.creds = &auth.Credentials{
		UserAuthentication: authentication,
		MantaKeyId:         "",
		MantaEndpoint:      auth.Endpoint{URL: s.Server.URL},
	}
	s.manta = localmanta.New(s.creds.MantaEndpoint.URL, s.creds.UserAuthentication.User)
	s.manta.SetupHTTP(s.Mux)
}

func (s *LocalTests) TearDownSuite(c *gc.C) {
	s.Mux = nil
	s.Server.Config.Handler = s.oldHandler
	s.Server.Close()
}

func (s *LocalTests) SetUpTest(c *gc.C) {
	client := client.NewClient(s.creds.MantaEndpoint.URL, "", s.creds, &manta.Logger)
	c.Assert(client, gc.NotNil)
	s.testClient = manta.New(client)
	c.Assert(s.testClient, gc.NotNil)
}

// Helper method to create a test directory
func (s *LocalTests) createDirectory(c *gc.C, path string) {
	err := s.testClient.PutDirectory(path)
	c.Assert(err, gc.IsNil)
}

// Helper method to create a test object
func (s *LocalTests) createObject(c *gc.C, path, objName string) {
	err := s.testClient.PutObject(path, objName, []byte("Test Manta API"))
	c.Assert(err, gc.IsNil)
}

// Helper method to delete a test directory
func (s *LocalTests) deleteDirectory(c *gc.C, path string) {
	err := s.testClient.DeleteDirectory(path)
	c.Assert(err, gc.IsNil)
}

// Helper method to delete a test object
func (s *LocalTests) deleteObject(c *gc.C, path, objName string) {
	err := s.testClient.DeleteObject(path, objName)
	c.Assert(err, gc.IsNil)
}

// Helper method to create a test job
func (s *LocalTests) createJob(c *gc.C, jobName string) string {
	phases := []manta.Phase{
		{Type: "map", Exec: "wc", Init: ""},
		{Type: "reduce", Exec: "awk '{ l += $1; w += $2; c += $3 } END { print l, w, c }'", Init: ""},
	}
	jobUri, err := s.testClient.CreateJob(manta.CreateJobOpts{Name: jobName, Phases: phases})
	c.Assert(err, gc.IsNil)
	c.Assert(jobUri, gc.NotNil)
	return strings.Split(jobUri, "/")[3]
}

// Storage API
func (s *LocalTests) TestPutDirectory(c *gc.C) {
	s.createDirectory(c, "test")

	// cleanup
	s.deleteDirectory(c, "test")
}

func (s *LocalTests) TestListDirectory(c *gc.C) {
	s.createDirectory(c, "test")
	defer s.deleteDirectory(c, "test")
	s.createObject(c, "test", "obj")
	defer s.deleteObject(c, "test", "obj")

	opts := manta.ListDirectoryOpts{}
	dirs, err := s.testClient.ListDirectory("test", opts)
	c.Assert(err, gc.IsNil)
	c.Assert(dirs, gc.NotNil)
}

func (s *LocalTests) TestDeleteDirectory(c *gc.C) {
	s.createDirectory(c, "test")

	s.deleteDirectory(c, "test")
}

func (s *LocalTests) TestPutObject(c *gc.C) {
	s.createDirectory(c, "dir")
	defer s.deleteDirectory(c, "dir")
	s.createObject(c, "dir", "obj")
	defer s.deleteObject(c, "dir", "obj")
}

func (s *LocalTests) TestGetObject(c *gc.C) {
	s.createDirectory(c, "dir")
	defer s.deleteDirectory(c, "dir")
	s.createObject(c, "dir", "obj")
	defer s.deleteObject(c, "dir", "obj")

	obj, err := s.testClient.GetObject("dir", "obj")
	c.Assert(err, gc.IsNil)
	c.Assert(obj, gc.NotNil)
	c.Check(string(obj), gc.Equals, "Test Manta API")
}

func (s *LocalTests) TestDeleteObject(c *gc.C) {
	s.createDirectory(c, "dir")
	defer s.deleteDirectory(c, "dir")
	s.createObject(c, "dir", "obj")

	s.deleteObject(c, "dir", "obj")
}

func (s *LocalTests) TestPutSnapLink(c *gc.C) {
	s.createDirectory(c, "linkdir")
	defer s.deleteDirectory(c, "linkdir")
	s.createObject(c, "linkdir", "obj")
	defer s.deleteObject(c, "linkdir", "obj")

	location := fmt.Sprintf("/%s/%s", s.creds.UserAuthentication.User, "stor/linkdir/obj")
	err := s.testClient.PutSnapLink("linkdir", "objlnk", location)
	c.Assert(err, gc.IsNil)

	// cleanup
	s.deleteObject(c, "linkdir", "objlnk")
}

// Jobs API
func (s *LocalTests) TestCreateJob(c *gc.C) {
	s.createJob(c, "test-job")
}

func (s *LocalTests) TestListLiveJobs(c *gc.C) {
	s.createJob(c, "test-job")

	jobs, err := s.testClient.ListJobs(true)
	c.Assert(err, gc.IsNil)
	c.Assert(jobs, gc.NotNil)
	c.Assert(len(jobs) >= 1, gc.Equals, true)
}

func (s *LocalTests) TestListAllJobs(c *gc.C) {
	s.createJob(c, "test-job")

	jobs, err := s.testClient.ListJobs(false)
	c.Assert(err, gc.IsNil)
	c.Assert(jobs, gc.NotNil)
}

func (s *LocalTests) TestCancelJob(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	err := s.testClient.CancelJob(jobId)
	c.Assert(err, gc.IsNil)
}

func (s *LocalTests) TestAddJobInputs(c *gc.C) {
	var inputs = `/` + s.creds.UserAuthentication.User + `/stor/testjob/obj1
/` + s.creds.UserAuthentication.User + `/stor/testjob/obj2
`
	s.createDirectory(c, "testjob")
	defer s.deleteDirectory(c, "testjob")
	s.createObject(c, "testjob", "obj1")
	defer s.deleteObject(c, "testjob", "obj1")
	s.createObject(c, "testjob", "obj2")
	defer s.deleteObject(c, "testjob", "obj2")
	jobId := s.createJob(c, "test-job")

	err := s.testClient.AddJobInputs(jobId, strings.NewReader(inputs))
	c.Assert(err, gc.IsNil)
}

func (s *LocalTests) TestEndJobInputs(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	err := s.testClient.EndJobInputs(jobId)
	c.Assert(err, gc.IsNil)
}

func (s *LocalTests) TestGetJob(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	job, err := s.testClient.GetJob(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(job, gc.NotNil)
}

func (s *LocalTests) TestGetJobInput(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	input, err := s.testClient.GetJobInput(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(input, gc.NotNil)
}

func (s *LocalTests) TestGetJobOutput(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	output, err := s.testClient.GetJobOutput(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(output, gc.NotNil)
}

func (s *LocalTests) TestGetJobFailures(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	fail, err := s.testClient.GetJobFailures(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(fail, gc.Equals, "")
}

func (s *LocalTests) TestGetJobErrors(c *gc.C) {
	jobId := s.createJob(c, "test-job")

	errs, err := s.testClient.GetJobErrors(jobId)
	c.Assert(err, gc.IsNil)
	c.Assert(errs, gc.IsNil)
}
<|endoftext|>


<commit_before>package twitch

import (
	"bufio"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/sorcix/irc"
)

// a message on the queue, this is not what the outside world sees
type queueItem struct {
	message irc.Message
	signal  chan bool
}

type TwitchClient struct {
	server   string
	username string
	password string

	// connection handling
	conn   net.Conn
	reader *bufio.Reader
	writer *irc.Encoder

	// handlers for incoming messages
	handlers map[string]HandlerFunc

	// time between two regular messages are sent
	delay time.Duration

	// this signal is sent when the client has sent the CAP REQ commands
	ready chan struct{}

	// this signal is sent when we disconnected
	alive chan struct{}

	// these are fired when .Disconnect() is called
	stopSending   chan struct{}
	stopReceiving chan struct{}

	// this is fired when .sender() / .receiver() stop
	stoppedSending   chan struct{}
	stoppedReceiving chan struct{}

	// on this channel incoming messages from the network are sent
	incoming chan Message

	// list of ougtoing messages (sent by us)
	outgoing chan queueItem
}

func NewTwitchClient(server string, username string, password string, delay time.Duration) *TwitchClient {
	client := &TwitchClient{
		server:           server,
		username:         username,
		password:         password,
		delay:            delay,
		conn:             nil,
		reader:           nil,
		writer:           nil,
		ready:            make(chan struct{}),
		alive:            make(chan struct{}),
		stopReceiving:    make(chan struct{}),
		stopSending:      make(chan struct{}),
		stoppedReceiving: make(chan struct{}),
		stoppedSending:   make(chan struct{}),
		incoming:         make(chan Message, 50),
		outgoing:         make(chan queueItem, 50),
	}

	// setup vital message listeners
	client.setupHandlers()

	return client
}

func (client *TwitchClient) Ready() <-chan struct{} {
	return client.ready
}

func (client *TwitchClient) Alive() <-chan struct{} {
	return client.alive
}

func (client *TwitchClient) Incoming() <-chan Message {
	return client.incoming
}

func (client *TwitchClient) Connect() error {
	conn, err := net.Dial("tcp", client.server)
	if err != nil {
		return err
	}

	client.conn = conn
	client.reader = bufio.NewReader(conn) // we manually read to properly handle tags
	client.writer = irc.NewEncoder(conn)

	// start working on the queue
	go client.sender()

	// start receiving
	go client.receiver()

	// send login info before anything else
	client.Send(irc.Message{
		Command: irc.PASS,
		Params:  []string{client.password},
	})

	client.Send(irc.Message{
		Command: irc.NICK,
		Params:  []string{client.username},
	})

	client.Send(irc.Message{
		Command: irc.USER,
		Params:  []string{"kabukibot", "8", "*", client.username},
	})

	return nil
}

func (client *TwitchClient) Disconnect() error {
	// stop the sender/receiver and wait for them to stop (maybe it will drain the
	// outgoing queue, maybe it won't, but let's give it time)
	close(client.stopReceiving)
	<-client.stoppedReceiving

	close(client.stopSending)
	<-client.stoppedSending

	// for all intents and purposes, we are not alive anymore
	close(client.alive)

	// close the IRC connection
	return client.conn.Close()
}

func (client *TwitchClient) Send(msg irc.Message) <-chan bool {
	signal := make(chan bool, 1)
	outgoing := queueItem{msg, signal}

	// if queue is not full, then
	client.outgoing <- outgoing
	// else
	// 	signal <- false // means "not sent"
	// 	close(signal)

	return signal
}

func (client *TwitchClient) sender() {
	for {
		select {
		case msg := <-client.outgoing:
			fmt.Println("< " + msg.message.String())
			client.writer.Encode(&msg.message)

			// signal to the one who sent the message that it was in fact sent
			msg.signal <- true
			close(msg.signal)

		case <-client.stopSending:
			break
		}
	}

	close(client.stoppedSending)
}

func (client *TwitchClient) receiver() {
	reading := make(chan struct{})

	// a buffer between the raw irc input from the net and the goroutine channels
	buffer := make(chan string, 10)

	// fork a reader loop, which could block and needs special handling as it's not a channel
	// (but it will pump its messages into a channel)
	go func() {
		for {
			select {
			case <-client.stopReceiving:
				break

			default:
				// set a 5min timeout
				client.conn.SetDeadline(time.Now().Add(300 * time.Second))

				line, err := client.reader.ReadString('\n')
				if err != nil {
					break
				}

				buffer <- line
			}
		}

		close(reading)
	}()

	for {
		select {
		case rawLine := <-buffer:
			fmt.Println("> " + strings.TrimSpace(rawLine))
			// if the message begins with a '@', we have some tags (IRCv3). The default
			// IRC decoder will not have properly detected it and mangled its output.
			// We fix that by manually splitting the tags from the rest of the message
			// and parse each part individually.
			tags := make(irc.Tags)
			msg := &irc.Message{}

			if strings.HasPrefix(rawLine, "@") {
				parts := strings.SplitN(rawLine, " ", 2)

				tags = irc.ParseTags(strings.TrimPrefix(parts[0], "@"))
				msg = irc.ParseMessage(parts[1])
			} else {
				msg = irc.ParseMessage(rawLine)
			}

			// hand it over to the message handler
			handler, ok := client.handlers[msg.Command]
			if ok {
				go handler(msg, tags)
			}

		case <-client.stopReceiving:
			break
		}
	}

	close(client.stoppedReceiving)
}
<commit_msg>limit number of outgoing messages in the queue to 50<commit_after>package twitch

import (
	"bufio"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/sorcix/irc"
)

// buffer at most this many messages before dropping messages
// (this applies to OUTGOING messages)
const queueSize = 50

// a message on the queue, this is not what the outside world sees
type queueItem struct {
	message irc.Message
	signal  chan bool
}

type TwitchClient struct {
	server   string
	username string
	password string

	// connection handling
	conn   net.Conn
	reader *bufio.Reader
	writer *irc.Encoder

	// handlers for incoming messages
	handlers map[string]HandlerFunc

	// time between two regular messages are sent
	delay time.Duration

	// this signal is sent when the client has sent the CAP REQ commands
	ready chan struct{}

	// this signal is sent when we disconnected
	alive chan struct{}

	// these are fired when .Disconnect() is called
	stopSending   chan struct{}
	stopReceiving chan struct{}

	// this is fired when .sender() / .receiver() stop
	stoppedSending   chan struct{}
	stoppedReceiving chan struct{}

	// on this channel incoming messages from the network are sent
	incoming chan Message

	// list of ougtoing messages (sent by us)
	outgoing   chan queueItem
	queueLen   int
	queueSize  int
	queueMutex sync.Mutex
}

func NewTwitchClient(server string, username string, password string, delay time.Duration) *TwitchClient {
	client := &TwitchClient{
		server:           server,
		username:         username,
		password:         password,
		delay:            delay,
		conn:             nil,
		reader:           nil,
		writer:           nil,
		ready:            make(chan struct{}),
		alive:            make(chan struct{}),
		stopReceiving:    make(chan struct{}),
		stopSending:      make(chan struct{}),
		stoppedReceiving: make(chan struct{}),
		stoppedSending:   make(chan struct{}),
		incoming:         make(chan Message, 50),
		outgoing:         make(chan queueItem, queueSize+10), // make a bit room so we never block, even when reaching queueSize
		queueSize:        queueSize,
		queueLen:         0,
		queueMutex:       sync.Mutex{},
	}

	// setup vital message listeners
	client.setupHandlers()

	return client
}

func (client *TwitchClient) Ready() <-chan struct{} {
	return client.ready
}

func (client *TwitchClient) Alive() <-chan struct{} {
	return client.alive
}

func (client *TwitchClient) Incoming() <-chan Message {
	return client.incoming
}

func (client *TwitchClient) Connect() error {
	conn, err := net.Dial("tcp", client.server)
	if err != nil {
		return err
	}

	client.conn = conn
	client.reader = bufio.NewReader(conn) // we manually read to properly handle tags
	client.writer = irc.NewEncoder(conn)

	// start working on the queue
	go client.sender()

	// start receiving
	go client.receiver()

	// send login info before anything else
	client.Send(irc.Message{
		Command: irc.PASS,
		Params:  []string{client.password},
	})

	client.Send(irc.Message{
		Command: irc.NICK,
		Params:  []string{client.username},
	})

	client.Send(irc.Message{
		Command: irc.USER,
		Params:  []string{"kabukibot", "8", "*", client.username},
	})

	return nil
}

func (client *TwitchClient) Disconnect() error {
	// stop the sender/receiver and wait for them to stop (maybe it will drain the
	// outgoing queue, maybe it won't, but let's give it time)
	close(client.stopReceiving)
	<-client.stoppedReceiving

	close(client.stopSending)
	<-client.stoppedSending

	// for all intents and purposes, we are not alive anymore
	close(client.alive)

	// close the IRC connection
	return client.conn.Close()
}

func (client *TwitchClient) Send(msg irc.Message) <-chan bool {
	signal := make(chan bool, 1)

	client.queueMutex.Lock()

	// silenty drop the message so our queue doesn't grow infinitely
	if client.queueLen >= client.queueSize {
		signal <- false
		close(signal)
	} else {
		client.outgoing <- queueItem{msg, signal}
		client.queueLen++
	}

	client.queueMutex.Unlock()

	return signal
}

func (client *TwitchClient) sender() {
	for {
		select {
		case msg := <-client.outgoing:
			fmt.Println("< " + msg.message.String())
			client.writer.Encode(&msg.message)

			// signal to the one who sent the message that it was in fact sent
			msg.signal <- true
			close(msg.signal)

			client.queueMutex.Lock()
			client.queueLen--
			client.queueMutex.Unlock()

		case <-client.stopSending:
			break
		}
	}

	close(client.stoppedSending)
}

func (client *TwitchClient) receiver() {
	reading := make(chan struct{})

	// a buffer between the raw irc input from the net and the goroutine channels
	buffer := make(chan string, 10)

	// fork a reader loop, which could block and needs special handling as it's not a channel
	// (but it will pump its messages into a channel)
	go func() {
		for {
			select {
			case <-client.stopReceiving:
				break

			default:
				// set a 5min timeout
				client.conn.SetDeadline(time.Now().Add(300 * time.Second))

				line, err := client.reader.ReadString('\n')
				if err != nil {
					break
				}

				buffer <- line
			}
		}

		close(reading)
	}()

	for {
		select {
		case rawLine := <-buffer:
			fmt.Println("> " + strings.TrimSpace(rawLine))
			// if the message begins with a '@', we have some tags (IRCv3). The default
			// IRC decoder will not have properly detected it and mangled its output.
			// We fix that by manually splitting the tags from the rest of the message
			// and parse each part individually.
			tags := make(irc.Tags)
			msg := &irc.Message{}

			if strings.HasPrefix(rawLine, "@") {
				parts := strings.SplitN(rawLine, " ", 2)

				tags = irc.ParseTags(strings.TrimPrefix(parts[0], "@"))
				msg = irc.ParseMessage(parts[1])
			} else {
				msg = irc.ParseMessage(rawLine)
			}

			// hand it over to the message handler
			handler, ok := client.handlers[msg.Command]
			if ok {
				go handler(msg, tags)
			}

		case <-client.stopReceiving:
			break
		}
	}

	close(client.stoppedReceiving)
}
<|endoftext|>


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

package upgrades

import (
	"fmt"

	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/state"
)

func updateRsyslogPort(context Context) error {
	agentConfig := context.AgentConfig()
	info, ok := agentConfig.StateInfo()
	if !ok {
		return fmt.Errorf("Failed to get StateInfo")
	}
	// we need to re-open state so we can bypass validation
	// as the syslog-port is normally immutable
	st, err := state.Open(info, state.DefaultDialOpts(), nil)
	if err != nil {
		return err
	}
	defer st.Close()
	attrs := map[string]interface{}{
		"syslog-port": config.DefaultSyslogPort,
	}
	return st.UpdateEnvironConfig(attrs, nil, nil)
}
<commit_msg>Improve comment<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package upgrades

import (
	"fmt"

	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/state"
)

func updateRsyslogPort(context Context) error {
	agentConfig := context.AgentConfig()
	info, ok := agentConfig.StateInfo()
	if !ok {
		return fmt.Errorf("Failed to get StateInfo")
	}
	// we need to re-open state with a nil policay so we can bypass
	// validation, as the syslog-port is normally immutable
	st, err := state.Open(info, state.DefaultDialOpts(), nil)
	if err != nil {
		return err
	}
	defer st.Close()
	attrs := map[string]interface{}{
		"syslog-port": config.DefaultSyslogPort,
	}
	return st.UpdateEnvironConfig(attrs, nil, nil)
}
<|endoftext|>


<commit_before>// +build !go1.7

package transport

import "crypto/tls"

// TLSConfigClone returns a clone of tls.Config. This function is provided for
// compatibility for go1.7 that doesn't include this method in stdlib.
func TLSConfigClone(c *tls.Config) *tls.Config {
	return c.Clone()
}
<commit_msg>update the transport layer to work with all crypto tls version<commit_after>// +build !go1.7

package transport

import "crypto/tls"

// TLSConfigClone returns a clone of tls.Config. This function is provided for
// compatibility for go1.7 that doesn't include this method in stdlib.
func TLSConfigClone(c *tls.Config) *tls.Config {
	return &tls.Config{
		Rand:  c.Rand,
		Time:  c.Time,
		Certificates:  c.Certificates,
		NameToCertificate:  c.NameToCertificate,
		GetCertificate:  c.GetCertificate,
		RootCAs:  c.RootCAs,
		NextProtos:  c.NextProtos,
		ServerName:  c.ServerName,
		ClientAuth:  c.ClientAuth,
		ClientCAs:  c.ClientCAs,
		InsecureSkipVerify:  c.InsecureSkipVerify,
		CipherSuites:  c.CipherSuites,
		PreferServerCipherSuites:  c.PreferServerCipherSuites,
		SessionTicketsDisabled:  c.SessionTicketsDisabled,
		SessionTicketKey:  c.SessionTicketKey,
		ClientSessionCache:  c.ClientSessionCache,
		MinVersion:  c.MinVersion,
		MaxVersion:  c.MaxVersion,
		CurvePreferences:  c.CurvePreferences,
	}
}
<|endoftext|>


<commit_before>package advhttp

import (
	"bufio"
	"errors"
	"net"
	"net/http"
	"path"
	"strings"
)

// The util response writer gives you the ability to hook into the response
// writer lifecycle. This Includes being able to have a callback to re-write
// the headers before the body is written out. This is mostly useful for a
// response writer to give to a reverse proxy. You can then re-write the
// location header as well as add in via headers, etc.
type UtilResponseWriter struct {
	w                   http.ResponseWriter
	sentHeaders         bool
	SendHeadersCallback func(http.Header)
}

// Creates a new UtilResponseWriter wrapping the given http.UtilResponseWriter
func NewUtilResponseWriter(w http.ResponseWriter, hcb func(http.Header)) *UtilResponseWriter {
	return &UtilResponseWriter{w: w, SendHeadersCallback: hcb}
}

// Creates a new UtilResponseWriter wrapping the given http.ResponseWriter with
// a location rewriting and via header adding response writer
func NewRPResponseWriter(w http.ResponseWriter, prefix string, via string) *UtilResponseWriter {
	return &UtilResponseWriter{w: w, SendHeadersCallback: NewReverseProxyHeadersCallback(prefix, via)}
}

func NewRewriteLocationHeaderCallback(prefix string) func(http.Header) {
	return func(headers http.Header) {
		oldLocation := headers.Get("Location")
		if oldLocation != "" {
			if oldLocation != "/" {
				if strings.HasPrefix(oldLocation, "/") {
					if strings.HasSuffix(oldLocation, "/") {
						headers.Set("Location", path.Join(prefix, oldLocation)+"/")
					} else {
						headers.Set("Location", path.Join(prefix, oldLocation))
					}
				}
			}
		}
	}
}

func NewReverseProxyHeadersCallback(prefix string, via string) func(http.Header) {
	return func(headers http.Header) {
		oldLocation := headers.Get("Location")
		if oldLocation != "" {
			if strings.HasPrefix(oldLocation, "/") {
				if strings.HasSuffix(oldLocation, "/") {
					headers.Set("Location", path.Join(prefix, oldLocation)+"/")
				} else {
					headers.Set("Location", path.Join(prefix, oldLocation))
				}
			}
		}
		if headers.Get("Via") == "" {
			headers.Set("Via", via)
		} else {
			headers.Set("Via", via+", "+headers.Get("Via"))
		}
	}
}

func (urw *UtilResponseWriter) Header() http.Header {
	return urw.w.Header()
}

func (urw *UtilResponseWriter) WriteHeader(status int) {
	urw.SendHeadersCallback(urw.Header())
	urw.sentHeaders = true
	urw.w.WriteHeader(status)
}

func (urw *UtilResponseWriter) Write(bytes []byte) (int, error) {
	if !urw.sentHeaders {
		urw.SendHeadersCallback(urw.Header())
		urw.sentHeaders = true
	}
	return urw.w.Write(bytes)
}

func (urw *UtilResponseWriter) GetFlusher() (flusher http.Flusher, ok bool) {
	flusher, ok = urw.w.(http.Flusher)
	return
}

func (urw *UtilResponseWriter) Flush() {
	if flusher, ok := urw.w.(http.Flusher); ok {
		flusher.Flush()
	}
}

func (urw *UtilResponseWriter) GetHijacker() (hijacker http.Hijacker, ok bool) {
	hijacker, ok = urw.w.(http.Hijacker)
	return
}

func (urw *UtilResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := urw.w.(http.Hijacker); ok {
		return hijacker.Hijack()
	}
	return nil, nil, errors.New("Couldn't cast responsewriter to hijacker")
}

func (urw *UtilResponseWriter) GetCloseNotifier() (closeNotifier http.CloseNotifier, ok bool) {
	closeNotifier, ok = urw.w.(http.CloseNotifier)
	return
}

func (urw *UtilResponseWriter) CloseNotify() <-chan bool {
	if closeNotifier, ok := urw.w.(http.CloseNotifier); ok {
		return closeNotifier.CloseNotify()
	}
	return nil
}
<commit_msg>Not for root path<commit_after>package advhttp

import (
	"bufio"
	"errors"
	"net"
	"net/http"
	"path"
	"strings"
)

// The util response writer gives you the ability to hook into the response
// writer lifecycle. This Includes being able to have a callback to re-write
// the headers before the body is written out. This is mostly useful for a
// response writer to give to a reverse proxy. You can then re-write the
// location header as well as add in via headers, etc.
type UtilResponseWriter struct {
	w                   http.ResponseWriter
	sentHeaders         bool
	SendHeadersCallback func(http.Header)
}

// Creates a new UtilResponseWriter wrapping the given http.UtilResponseWriter
func NewUtilResponseWriter(w http.ResponseWriter, hcb func(http.Header)) *UtilResponseWriter {
	return &UtilResponseWriter{w: w, SendHeadersCallback: hcb}
}

// Creates a new UtilResponseWriter wrapping the given http.ResponseWriter with
// a location rewriting and via header adding response writer
func NewRPResponseWriter(w http.ResponseWriter, prefix string, via string) *UtilResponseWriter {
	return &UtilResponseWriter{w: w, SendHeadersCallback: NewReverseProxyHeadersCallback(prefix, via)}
}

func NewRewriteLocationHeaderCallback(prefix string) func(http.Header) {
	return func(headers http.Header) {
		oldLocation := headers.Get("Location")
		if oldLocation != "" {
			if strings.HasPrefix(oldLocation, "/") {
				if strings.HasSuffix(oldLocation, "/") && oldLocation != "/" {
					headers.Set("Location", path.Join(prefix, oldLocation)+"/")
				} else {
					headers.Set("Location", path.Join(prefix, oldLocation))
				}
			}
		}
	}
}

func NewReverseProxyHeadersCallback(prefix string, via string) func(http.Header) {
	return func(headers http.Header) {
		oldLocation := headers.Get("Location")
		if oldLocation != "" {
			if strings.HasPrefix(oldLocation, "/") {
				if strings.HasSuffix(oldLocation, "/") && oldLocation != "/" {
					headers.Set("Location", path.Join(prefix, oldLocation)+"/")
				} else {
					headers.Set("Location", path.Join(prefix, oldLocation))
				}
			}
		}
		if headers.Get("Via") == "" {
			headers.Set("Via", via)
		} else {
			headers.Set("Via", via+", "+headers.Get("Via"))
		}
	}
}

func (urw *UtilResponseWriter) Header() http.Header {
	return urw.w.Header()
}

func (urw *UtilResponseWriter) WriteHeader(status int) {
	urw.SendHeadersCallback(urw.Header())
	urw.sentHeaders = true
	urw.w.WriteHeader(status)
}

func (urw *UtilResponseWriter) Write(bytes []byte) (int, error) {
	if !urw.sentHeaders {
		urw.SendHeadersCallback(urw.Header())
		urw.sentHeaders = true
	}
	return urw.w.Write(bytes)
}

func (urw *UtilResponseWriter) GetFlusher() (flusher http.Flusher, ok bool) {
	flusher, ok = urw.w.(http.Flusher)
	return
}

func (urw *UtilResponseWriter) Flush() {
	if flusher, ok := urw.w.(http.Flusher); ok {
		flusher.Flush()
	}
}

func (urw *UtilResponseWriter) GetHijacker() (hijacker http.Hijacker, ok bool) {
	hijacker, ok = urw.w.(http.Hijacker)
	return
}

func (urw *UtilResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := urw.w.(http.Hijacker); ok {
		return hijacker.Hijack()
	}
	return nil, nil, errors.New("Couldn't cast responsewriter to hijacker")
}

func (urw *UtilResponseWriter) GetCloseNotifier() (closeNotifier http.CloseNotifier, ok bool) {
	closeNotifier, ok = urw.w.(http.CloseNotifier)
	return
}

func (urw *UtilResponseWriter) CloseNotify() <-chan bool {
	if closeNotifier, ok := urw.w.(http.CloseNotifier); ok {
		return closeNotifier.CloseNotify()
	}
	return nil
}
<|endoftext|>


<commit_before>package watcher

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"
)

// setup creates all required files and folders for
// the tests and returns a function that is used as
// a teardown function when the tests are done.
func setup(t *testing.T) (string, func()) {
	testDir, err := ioutil.TempDir(".", "")
	if err != nil {
		t.Error(err)
	}

	err = ioutil.WriteFile(filepath.Join(testDir, "file.txt"),
		[]byte{}, 0755)
	if err != nil {
		t.Error(err)
	}

	testDirTwo := filepath.Join(testDir, "testDirTwo")
	err = os.Mkdir(testDirTwo, 0755)
	if err != nil {
		t.Error(err)
	}

	err = ioutil.WriteFile(filepath.Join(testDirTwo, "file_recursive.txt"),
		[]byte{}, 0755)
	if err != nil {
		t.Error(err)
	}

	return testDir, func() {
		if os.RemoveAll(testDir); err != nil {
			t.Error(err)
		}
	}
}

func TestWatcherAdd(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	// Make sure len(w.Files) is 4.
	if len(w.Files) != 4 {
		t.Errorf("expected 4 files, found %d", len(w.Files))
	}

	// Make sure w.Names[0] is now equal to testDir.
	if w.Names[0] != testDir {
		t.Errorf("expected w.Names[0] to be %s, got %s",
			testDir, w.Names[0])
	}

	if _, found := w.Files["testDirTwo"]; !found {
		t.Error("expected to find testDirTwo directory")
	}

	if w.Files["testDirTwo"].Name() != "testDirTwo" {
		t.Errorf("expected w.Files[%q].Name() to be testDirTwo, got %s",
			"testDirTwo", w.Files["testDirTwo"].Name())
	}

	if _, found := w.Files["testDirTwo/file_recursive.txt"]; !found {
		t.Error("expected to find testDirTwo/file_recursive.txt directory")
	}

	if w.Files["testDirTwo/file_recursive.txt"].Name() != "file_recursive.txt" {
		t.Errorf("expected w.Files[%q].Name() to be file_recursive.txt, got %s",
			"testDirTwo/file_recursive.txt", w.Files["testDirTwo/file_recursive.txt"].Name())
	}
}

func TestWatcherAddNotFound(t *testing.T) {
	w := New()

	// Make sure there is an error when adding a
	// non-existent file/folder.
	if err := w.Add("random_filename.txt"); err == nil {
		t.Error("expected a file not found error")
	}
}

func TestWatcherRemove(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	// Add the testDir to the watchlist.
	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	// Make sure len(w.Files) is 4.
	if len(w.Files) != 4 {
		t.Errorf("expected 4 files, found %d", len(w.Files))
	}

	// Now remove the folder from the watchlist.
	if err := w.Remove(testDir); err != nil {
		t.Error(err)
	}

	// Now check that there is nothing being watched.
	if len(w.Files) != 0 {
		t.Errorf("expected len(w.Files) to be 0, got %d", len(w.Files))
	}

	// Make sure len(w.Names) is now 0.
	if len(w.Names) != 0 {
		t.Errorf("expected len(w.Names) to be empty, len(w.Names): %s", len(w.Names))
	}
}

func TestListFiles(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	fileList, err := ListFiles(testDir)
	if err != nil {
		t.Error(err)
	}

	// Make sure fInfoTest contains the correct os.FileInfo names.
	fname := filepath.Join(testDir, "file.txt")
	if fileList[fname].Name() != "file.txt" {
		t.Errorf("expected fileList[%s].Name() to be file.txt, got %s",
			fname, fileList[fname].Name())
	}
}

func TestTriggerEvent(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	// Add the testDir to the watchlist.
	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		select {
		case event := <-w.Event:
			if event.Name() != "triggered event" {
				t.Errorf("expected event file name to be triggered event, got %s",
					event.Name())
			}
			wg.Done()
		case <-time.After(time.Millisecond * 250):
			t.Error("received no event from Event channel")
			wg.Done()
		}
	}()

	go func() {
		// Start the watching process.
		if err := w.Start(time.Millisecond * 100); err != nil {
			t.Error(err)
		}
	}()

	w.TriggerEvent(EventFileAdded, nil)

	wg.Wait()
}

func TestEventAddFile(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	// Add the testDir to the watchlist.
	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	newFileName := filepath.Join(testDir, "newfile.txt")
	err := ioutil.WriteFile(newFileName, []byte{}, 0755)
	if err != nil {
		t.Error(err)
	}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		select {
		case event := <-w.Event:
			if event.EventType != EventFileAdded {
				t.Errorf("expected event to be EventFileAdded, got %s",
					event.EventType)
			}
			if event.Name() != "newfile.txt" {
				t.Errorf("expected event file name to be newfile.txt, got %s",
					event.Name())
			}
			wg.Done()
		case <-time.After(time.Millisecond * 250):
			t.Error("received no event from Event channel")
			wg.Done()
		}
	}()

	go func() {
		// Start the watching process.
		if err := w.Start(time.Millisecond * 100); err != nil {
			t.Error(err)
		}
	}()

	wg.Wait()
}

func TestEventDeleteFile(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	// Add the testDir to the watchlist.
	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		select {
		case <-w.Event:
			wg.Done()
		case <-time.After(time.Millisecond * 250):
			t.Error("received no event from Event channel")
			wg.Done()
		}
	}()

	go func() {
		// Start the watching process.
		if err := w.Start(time.Millisecond * 100); err != nil {
			t.Error(err)
		}
	}()

	if err := os.Remove(filepath.Join(testDir, "file.txt")); err != nil {
		t.Error(err)
	}

	wg.Wait()
}
<commit_msg>fixed TestTriggerEvent -race issue<commit_after>package watcher

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"
)

// setup creates all required files and folders for
// the tests and returns a function that is used as
// a teardown function when the tests are done.
func setup(t *testing.T) (string, func()) {
	testDir, err := ioutil.TempDir(".", "")
	if err != nil {
		t.Error(err)
	}

	err = ioutil.WriteFile(filepath.Join(testDir, "file.txt"),
		[]byte{}, 0755)
	if err != nil {
		t.Error(err)
	}

	testDirTwo := filepath.Join(testDir, "testDirTwo")
	err = os.Mkdir(testDirTwo, 0755)
	if err != nil {
		t.Error(err)
	}

	err = ioutil.WriteFile(filepath.Join(testDirTwo, "file_recursive.txt"),
		[]byte{}, 0755)
	if err != nil {
		t.Error(err)
	}

	return testDir, func() {
		if os.RemoveAll(testDir); err != nil {
			t.Error(err)
		}
	}
}

func TestWatcherAdd(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	// Make sure len(w.Files) is 4.
	if len(w.Files) != 4 {
		t.Errorf("expected 4 files, found %d", len(w.Files))
	}

	// Make sure w.Names[0] is now equal to testDir.
	if w.Names[0] != testDir {
		t.Errorf("expected w.Names[0] to be %s, got %s",
			testDir, w.Names[0])
	}

	if _, found := w.Files["testDirTwo"]; !found {
		t.Error("expected to find testDirTwo directory")
	}

	if w.Files["testDirTwo"].Name() != "testDirTwo" {
		t.Errorf("expected w.Files[%q].Name() to be testDirTwo, got %s",
			"testDirTwo", w.Files["testDirTwo"].Name())
	}

	if _, found := w.Files["testDirTwo/file_recursive.txt"]; !found {
		t.Error("expected to find testDirTwo/file_recursive.txt directory")
	}

	if w.Files["testDirTwo/file_recursive.txt"].Name() != "file_recursive.txt" {
		t.Errorf("expected w.Files[%q].Name() to be file_recursive.txt, got %s",
			"testDirTwo/file_recursive.txt", w.Files["testDirTwo/file_recursive.txt"].Name())
	}
}

func TestWatcherAddNotFound(t *testing.T) {
	w := New()

	// Make sure there is an error when adding a
	// non-existent file/folder.
	if err := w.Add("random_filename.txt"); err == nil {
		t.Error("expected a file not found error")
	}
}

func TestWatcherRemove(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	// Add the testDir to the watchlist.
	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	// Make sure len(w.Files) is 4.
	if len(w.Files) != 4 {
		t.Errorf("expected 4 files, found %d", len(w.Files))
	}

	// Now remove the folder from the watchlist.
	if err := w.Remove(testDir); err != nil {
		t.Error(err)
	}

	// Now check that there is nothing being watched.
	if len(w.Files) != 0 {
		t.Errorf("expected len(w.Files) to be 0, got %d", len(w.Files))
	}

	// Make sure len(w.Names) is now 0.
	if len(w.Names) != 0 {
		t.Errorf("expected len(w.Names) to be empty, len(w.Names): %s", len(w.Names))
	}
}

func TestListFiles(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	fileList, err := ListFiles(testDir)
	if err != nil {
		t.Error(err)
	}

	// Make sure fInfoTest contains the correct os.FileInfo names.
	fname := filepath.Join(testDir, "file.txt")
	if fileList[fname].Name() != "file.txt" {
		t.Errorf("expected fileList[%s].Name() to be file.txt, got %s",
			fname, fileList[fname].Name())
	}
}

func TestTriggerEvent(t *testing.T) {
	w := New()

	// Add the testDir to the watchlist.
	if err := w.Add("example/test_folder"); err != nil {
		t.Error(err)
	}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		select {
		case event := <-w.Event:
			if event.Name() != "triggered event" {
				t.Errorf("expected event file name to be triggered event, got %s",
					event.Name())
			}
			wg.Done()
		case <-time.After(time.Millisecond * 250):
			t.Error("received no event from Event channel")
			wg.Done()
		}
	}()

	go func() {
		// Start the watching process.
		if err := w.Start(time.Millisecond * 100); err != nil {
			t.Error(err)
		}
	}()

	w.TriggerEvent(EventFileAdded, nil)

	wg.Wait()
}

func TestEventAddFile(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	// Add the testDir to the watchlist.
	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	newFileName := filepath.Join(testDir, "newfile.txt")
	err := ioutil.WriteFile(newFileName, []byte{}, 0755)
	if err != nil {
		t.Error(err)
	}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		select {
		case event := <-w.Event:
			if event.EventType != EventFileAdded {
				t.Errorf("expected event to be EventFileAdded, got %s",
					event.EventType)
			}
			if event.Name() != "newfile.txt" {
				t.Errorf("expected event file name to be newfile.txt, got %s",
					event.Name())
			}
			wg.Done()
		case <-time.After(time.Millisecond * 250):
			t.Error("received no event from Event channel")
			wg.Done()
		}
	}()

	go func() {
		// Start the watching process.
		if err := w.Start(time.Millisecond * 100); err != nil {
			t.Error(err)
		}
	}()

	wg.Wait()
}

func TestEventDeleteFile(t *testing.T) {
	testDir, teardown := setup(t)
	defer teardown()

	w := New()

	// Add the testDir to the watchlist.
	if err := w.Add(testDir); err != nil {
		t.Error(err)
	}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		select {
		case <-w.Event:
			wg.Done()
		case <-time.After(time.Millisecond * 250):
			t.Error("received no event from Event channel")
			wg.Done()
		}
	}()

	go func() {
		// Start the watching process.
		if err := w.Start(time.Millisecond * 100); err != nil {
			t.Error(err)
		}
	}()

	if err := os.Remove(filepath.Join(testDir, "file.txt")); err != nil {
		t.Error(err)
	}

	wg.Wait()
}
<|endoftext|>