<commit_before>package luddite

import (
	"errors"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path"
	"runtime"
	"strings"
	"syscall"

	"github.com/SpirentOrion/httprouter"
	log "github.com/SpirentOrion/logrus"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/net/context"
)

const (
	DEFAULT_METRICS_URI_PATH = "/metrics"
)

// Service is an interface that implements a standalone RESTful web service.
type Service interface {
	// AddHandler adds a context-aware middleware handler to the
	// middleware stack. All handlers must be added before Run is
	// called.
	AddHandler(h Handler)

	// AddSingletonResource registers a singleton-style resource
	// (supporting GET and PUT methods only).
	AddSingletonResource(itemPath string, r Resource)

	// AddCollectionResource registers a collection-style resource
	// (supporting GET, POST, PUT, and DELETE methods).
	AddCollectionResource(basePath string, r Resource)

	// Config returns the service's ServiceConfig instance.
	Config() *ServiceConfig

	// Logger returns the service's log.Logger instance.
	Logger() *log.Logger

	// Router returns the service's httprouter.Router instance.
	Router() *httprouter.Router

	// Run is a convenience function that runs the service as an
	// HTTP server. The address is taken from the ServiceConfig
	// passed to NewService.
	Run() error
}

type service struct {
	config        *ServiceConfig
	defaultLogger *log.Logger
	accessLogger  *log.Logger
	router        *httprouter.Router
	handlers      []Handler
	middleware    *middleware
	schema        *SchemaHandler
}

// Verify that service implements Service.
var _ Service = &service{}

func NewService(config *ServiceConfig) (Service, error) {
	var err error

	// Create the service
	s := &service{
		config: config,
		router: httprouter.New(),
	}

	s.defaultLogger = log.New()
	if config.Log.ServiceLogPath != "" {
		openLogFile(s.defaultLogger, config.Log.ServiceLogPath)
		s.defaultLogger.Formatter = &log.JSONFormatter{}
	} else {
		s.defaultLogger.Out = os.Stdout
	}

	switch strings.ToLower(config.Log.ServiceLogLevel) {
	case "debug":
		s.defaultLogger.Level = log.DebugLevel
	default:
		fallthrough
	case "info":
		s.defaultLogger.Level = log.InfoLevel
	case "warn":
		s.defaultLogger.Level = log.WarnLevel
	case "error":
		s.defaultLogger.Level = log.ErrorLevel
	}

	// Add handler to log stacktrace
	addStackTraceHandler(s.defaultLogger)

	s.accessLogger = log.New()
	if config.Log.AccessLogPath != "" {
		openLogFile(s.accessLogger, config.Log.AccessLogPath)
		s.accessLogger.Formatter = &log.JSONFormatter{}
	} else {
		s.accessLogger.Out = os.Stdout
		s.accessLogger.Level = log.DebugLevel
	}

	s.router.NotFound = func(_ context.Context, rw http.ResponseWriter, _ *http.Request) {
		rw.WriteHeader(http.StatusNotFound)
	}

	s.router.MethodNotAllowed = func(_ context.Context, rw http.ResponseWriter, _ *http.Request) {
		rw.WriteHeader(http.StatusMethodNotAllowed)
	}

	// Create default middleware handlers
	bottom, err := s.newBottomHandler()
	if err != nil {
		return nil, err
	}

	negotiator, err := s.newNegotiatorHandler()
	if err != nil {
		return nil, err
	}

	context, err := s.newContextHandler()
	if err != nil {
		return nil, err
	}

	// Build middleware stack
	s.handlers = []Handler{bottom, negotiator, context}
	s.middleware = buildMiddleware(s.handlers)

	// Install default http handlers
	if s.config.Metrics.Enabled {
		s.addMetricsRoute()
	}
	if config.Schema.Enabled {
		s.addSchemaRoutes()
	}

	return s, nil
}

func (s *service) AddHandler(h Handler) {
	s.handlers = append(s.handlers, h)
	s.middleware = buildMiddleware(s.handlers)
}

func (s *service) AddSingletonResource(basePath string, r Resource) {
	// GET /basePath
	AddGetRoute(s.router, basePath, false, r)

	// PUT /basePath
	AddUpdateRoute(s.router, basePath, false, r)

	// POST /basePath/{action}
	AddActionRoute(s.router, basePath, false, r)
}

func (s *service) AddCollectionResource(basePath string, r Resource) {
	// GET /basePath
	AddListRoute(s.router, basePath, r)

	// GET /basePath/{id}
	AddGetRoute(s.router, basePath, true, r)

	// POST /basePath
	AddCreateRoute(s.router, basePath, r)

	// PUT /basePath/{id}
	AddUpdateRoute(s.router, basePath, true, r)

	// DELETE /basePath
	AddDeleteRoute(s.router, basePath, false, r)

	// DELETE /basePath/{id}
	AddDeleteRoute(s.router, basePath, true, r)

	// POST /basePath/{id}/{action}
	AddActionRoute(s.router, basePath, true, r)
}

func (s *service) Config() *ServiceConfig {
	return s.config
}

func (s *service) Logger() *log.Logger {
	return s.defaultLogger
}

func (s *service) Router() *httprouter.Router {
	return s.router
}

func (s *service) Run() error {
	// Add the router as the final middleware handler
	h, err := s.newRouterHandler()
	if err != nil {
		return err
	}
	s.AddHandler(h)

	var middleware http.Handler = s.middleware
	if s.config.Metrics.Enabled {
		middleware = prometheus.InstrumentHandler("service", middleware)
	}

	// Serve HTTP or HTTPS, depending on config. Use stoppable listener
	// so we can exit gracefully if signaled to do so.
	var stoppableListener net.Listener
	if s.config.Transport.TLS {
		s.defaultLogger.Debugf("HTTPS listening on %s", s.config.Addr)
		stoppableListener, err = NewStoppableTLSListener(s.config.Addr, true, s.config.Transport.CertFilePath, s.config.Transport.KeyFilePath)
	} else {
		s.defaultLogger.Debugf("HTTP listening on %s", s.config.Addr)
		stoppableListener, err = NewStoppableTCPListener(s.config.Addr, true)
	}
	if err != nil {
		return err
	}
	err = http.Serve(stoppableListener, middleware)
	if _, ok := err.(*ListenerStoppedError); ok {
		return nil
	}
	return err

}

func (s *service) newBottomHandler() (Handler, error) {
	return NewBottom(s.config, s.defaultLogger, s.accessLogger), nil
}

func (s *service) newNegotiatorHandler() (Handler, error) {
	return NewNegotiator([]string{ContentTypeJson, ContentTypeXml, ContentTypeHtml, ContentTypeOctetStream}), nil
}

func (s *service) newContextHandler() (Handler, error) {
	if s.config.Version.Min < 1 {
		return nil, errors.New("service's minimum API version must be greater than zero")
	}
	if s.config.Version.Max < 1 {
		return nil, errors.New("service's maximum API version must be greater than zero")
	}

	return NewContext(s, s.config.Version.Min, s.config.Version.Max), nil
}

func (s *service) newRouterHandler() (Handler, error) {
	return HandlerFunc(func(ctx context.Context, rw http.ResponseWriter, r *http.Request, _ ContextHandlerFunc) {
		// No more middleware handlers: further dispatch happens via httprouter
		s.router.HandleHTTP(ctx, rw, r)
	}), nil
}

func (s *service) addMetricsRoute() {
	uriPath := s.config.Metrics.UriPath
	if uriPath == "" {
		uriPath = DEFAULT_METRICS_URI_PATH
	}

	h := prometheus.UninstrumentedHandler()
	s.router.GET(uriPath, func(_ context.Context, rw http.ResponseWriter, r *http.Request) { h.ServeHTTP(rw, r) })
}

func (s *service) addSchemaRoutes() {
	config := s.config

	// Serve the various schemas, e.g. /schema/v1, /schema/v2, etc.
	s.schema = NewSchemaHandler(config.Schema.FilePath, config.Schema.FilePattern)
	s.router.GET(path.Join(config.Schema.UriPath, "/v:version"), s.schema.ServeHTTP)

	// Temporarily redirect (307) the base schema path to the default schema, e.g. /schema -> /schema/v2
	defaultSchemaPath := path.Join(config.Schema.UriPath, fmt.Sprintf("v%d", config.Version.Max))
	s.router.GET(config.Schema.UriPath, func(_ context.Context, rw http.ResponseWriter, r *http.Request) {
		http.Redirect(rw, r, defaultSchemaPath, http.StatusTemporaryRedirect)
	})

	// Optionally temporarily redirect (307) the root to the base schema path, e.g. / -> /schema
	if config.Schema.RootRedirect {
		s.router.GET("/", func(_ context.Context, rw http.ResponseWriter, r *http.Request) {
			http.Redirect(rw, r, config.Schema.UriPath, http.StatusTemporaryRedirect)
		})
	}
}

func openLogFile(logger *log.Logger, logPath string) {
	sigs := make(chan os.Signal, 1)
	logging := make(chan bool, 1)

	go func() {
		var curLog, priorLog *os.File
		for {
			// Open and begin using a new log file
			curLog, _ = os.OpenFile(logPath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0644)
			logger.SetOutput(curLog)

			if priorLog == nil {
				// First log, signal the outer goroutine that we're running
				logging <- true
			} else {
				// Follow-on log, close the prior log file
				priorLog.Close()
				priorLog = nil
			}

			// Wait for a SIGHUP
			<-sigs

			// Setup for the next iteration
			priorLog = curLog
		}
	}()

	signal.Notify(sigs, syscall.SIGHUP)
	<-logging
}

func addStackTraceHandler(logger *log.Logger) {
	sigs := make(chan os.Signal, 1)
	go func() {
		for {
			<-sigs
			buf := make([]byte, 1<<16)
			size := runtime.Stack(buf, true)
			logger.Infof("*** goroutine dump ***\n%s", buf[:size])
		}
	}()
	signal.Notify(sigs, syscall.SIGUSR1)
}
<commit_msg>Use JSONFormatter for default and access logging, including stdout<commit_after>package luddite

import (
	"errors"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path"
	"runtime"
	"strings"
	"syscall"

	"github.com/SpirentOrion/httprouter"
	log "github.com/SpirentOrion/logrus"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/net/context"
)

const (
	DEFAULT_METRICS_URI_PATH = "/metrics"
)

// Service is an interface that implements a standalone RESTful web service.
type Service interface {
	// AddHandler adds a context-aware middleware handler to the
	// middleware stack. All handlers must be added before Run is
	// called.
	AddHandler(h Handler)

	// AddSingletonResource registers a singleton-style resource
	// (supporting GET and PUT methods only).
	AddSingletonResource(itemPath string, r Resource)

	// AddCollectionResource registers a collection-style resource
	// (supporting GET, POST, PUT, and DELETE methods).
	AddCollectionResource(basePath string, r Resource)

	// Config returns the service's ServiceConfig instance.
	Config() *ServiceConfig

	// Logger returns the service's log.Logger instance.
	Logger() *log.Logger

	// Router returns the service's httprouter.Router instance.
	Router() *httprouter.Router

	// Run is a convenience function that runs the service as an
	// HTTP server. The address is taken from the ServiceConfig
	// passed to NewService.
	Run() error
}

type service struct {
	config        *ServiceConfig
	defaultLogger *log.Logger
	accessLogger  *log.Logger
	router        *httprouter.Router
	handlers      []Handler
	middleware    *middleware
	schema        *SchemaHandler
}

// Verify that service implements Service.
var _ Service = &service{}

func NewService(config *ServiceConfig) (Service, error) {
	var err error

	// Create the service
	s := &service{
		config: config,
		router: httprouter.New(),
	}

	s.defaultLogger = log.New()
	s.defaultLogger.SetFormatter(&log.JSONFormatter{})
	if config.Log.ServiceLogPath != "" {
		openLogFile(s.defaultLogger, config.Log.ServiceLogPath)
	} else {
		s.defaultLogger.SetOutput(os.Stdout)
	}

	switch strings.ToLower(config.Log.ServiceLogLevel) {
	case "debug":
		s.defaultLogger.SetLevel(log.DebugLevel)
	default:
		fallthrough
	case "info":
		s.defaultLogger.SetLevel(log.InfoLevel)
	case "warn":
		s.defaultLogger.SetLevel(log.WarnLevel)
	case "error":
		s.defaultLogger.SetLevel(log.ErrorLevel)
	}

	// Add handler to log stacktrace
	addStackTraceHandler(s.defaultLogger)

	s.accessLogger = log.New()
	s.accessLogger.SetFormatter(&log.JSONFormatter{})
	if config.Log.AccessLogPath != "" {
		openLogFile(s.accessLogger, config.Log.AccessLogPath)
	} else {
		s.accessLogger.SetOutput(os.Stdout)
		s.accessLogger.SetLevel(log.DebugLevel)
	}

	s.router.NotFound = func(_ context.Context, rw http.ResponseWriter, _ *http.Request) {
		rw.WriteHeader(http.StatusNotFound)
	}

	s.router.MethodNotAllowed = func(_ context.Context, rw http.ResponseWriter, _ *http.Request) {
		rw.WriteHeader(http.StatusMethodNotAllowed)
	}

	// Create default middleware handlers
	bottom, err := s.newBottomHandler()
	if err != nil {
		return nil, err
	}

	negotiator, err := s.newNegotiatorHandler()
	if err != nil {
		return nil, err
	}

	context, err := s.newContextHandler()
	if err != nil {
		return nil, err
	}

	// Build middleware stack
	s.handlers = []Handler{bottom, negotiator, context}
	s.middleware = buildMiddleware(s.handlers)

	// Install default http handlers
	if s.config.Metrics.Enabled {
		s.addMetricsRoute()
	}
	if config.Schema.Enabled {
		s.addSchemaRoutes()
	}

	return s, nil
}

func (s *service) AddHandler(h Handler) {
	s.handlers = append(s.handlers, h)
	s.middleware = buildMiddleware(s.handlers)
}

func (s *service) AddSingletonResource(basePath string, r Resource) {
	// GET /basePath
	AddGetRoute(s.router, basePath, false, r)

	// PUT /basePath
	AddUpdateRoute(s.router, basePath, false, r)

	// POST /basePath/{action}
	AddActionRoute(s.router, basePath, false, r)
}

func (s *service) AddCollectionResource(basePath string, r Resource) {
	// GET /basePath
	AddListRoute(s.router, basePath, r)

	// GET /basePath/{id}
	AddGetRoute(s.router, basePath, true, r)

	// POST /basePath
	AddCreateRoute(s.router, basePath, r)

	// PUT /basePath/{id}
	AddUpdateRoute(s.router, basePath, true, r)

	// DELETE /basePath
	AddDeleteRoute(s.router, basePath, false, r)

	// DELETE /basePath/{id}
	AddDeleteRoute(s.router, basePath, true, r)

	// POST /basePath/{id}/{action}
	AddActionRoute(s.router, basePath, true, r)
}

func (s *service) Config() *ServiceConfig {
	return s.config
}

func (s *service) Logger() *log.Logger {
	return s.defaultLogger
}

func (s *service) Router() *httprouter.Router {
	return s.router
}

func (s *service) Run() error {
	// Add the router as the final middleware handler
	h, err := s.newRouterHandler()
	if err != nil {
		return err
	}
	s.AddHandler(h)

	var middleware http.Handler = s.middleware
	if s.config.Metrics.Enabled {
		middleware = prometheus.InstrumentHandler("service", middleware)
	}

	// Serve HTTP or HTTPS, depending on config. Use stoppable listener
	// so we can exit gracefully if signaled to do so.
	var stoppableListener net.Listener
	if s.config.Transport.TLS {
		s.defaultLogger.Debugf("HTTPS listening on %s", s.config.Addr)
		stoppableListener, err = NewStoppableTLSListener(s.config.Addr, true, s.config.Transport.CertFilePath, s.config.Transport.KeyFilePath)
	} else {
		s.defaultLogger.Debugf("HTTP listening on %s", s.config.Addr)
		stoppableListener, err = NewStoppableTCPListener(s.config.Addr, true)
	}
	if err != nil {
		return err
	}
	err = http.Serve(stoppableListener, middleware)
	if _, ok := err.(*ListenerStoppedError); ok {
		return nil
	}
	return err

}

func (s *service) newBottomHandler() (Handler, error) {
	return NewBottom(s.config, s.defaultLogger, s.accessLogger), nil
}

func (s *service) newNegotiatorHandler() (Handler, error) {
	return NewNegotiator([]string{ContentTypeJson, ContentTypeXml, ContentTypeHtml, ContentTypeOctetStream}), nil
}

func (s *service) newContextHandler() (Handler, error) {
	if s.config.Version.Min < 1 {
		return nil, errors.New("service's minimum API version must be greater than zero")
	}
	if s.config.Version.Max < 1 {
		return nil, errors.New("service's maximum API version must be greater than zero")
	}

	return NewContext(s, s.config.Version.Min, s.config.Version.Max), nil
}

func (s *service) newRouterHandler() (Handler, error) {
	return HandlerFunc(func(ctx context.Context, rw http.ResponseWriter, r *http.Request, _ ContextHandlerFunc) {
		// No more middleware handlers: further dispatch happens via httprouter
		s.router.HandleHTTP(ctx, rw, r)
	}), nil
}

func (s *service) addMetricsRoute() {
	uriPath := s.config.Metrics.UriPath
	if uriPath == "" {
		uriPath = DEFAULT_METRICS_URI_PATH
	}

	h := prometheus.UninstrumentedHandler()
	s.router.GET(uriPath, func(_ context.Context, rw http.ResponseWriter, r *http.Request) { h.ServeHTTP(rw, r) })
}

func (s *service) addSchemaRoutes() {
	config := s.config

	// Serve the various schemas, e.g. /schema/v1, /schema/v2, etc.
	s.schema = NewSchemaHandler(config.Schema.FilePath, config.Schema.FilePattern)
	s.router.GET(path.Join(config.Schema.UriPath, "/v:version"), s.schema.ServeHTTP)

	// Temporarily redirect (307) the base schema path to the default schema, e.g. /schema -> /schema/v2
	defaultSchemaPath := path.Join(config.Schema.UriPath, fmt.Sprintf("v%d", config.Version.Max))
	s.router.GET(config.Schema.UriPath, func(_ context.Context, rw http.ResponseWriter, r *http.Request) {
		http.Redirect(rw, r, defaultSchemaPath, http.StatusTemporaryRedirect)
	})

	// Optionally temporarily redirect (307) the root to the base schema path, e.g. / -> /schema
	if config.Schema.RootRedirect {
		s.router.GET("/", func(_ context.Context, rw http.ResponseWriter, r *http.Request) {
			http.Redirect(rw, r, config.Schema.UriPath, http.StatusTemporaryRedirect)
		})
	}
}

func openLogFile(logger *log.Logger, logPath string) {
	sigs := make(chan os.Signal, 1)
	logging := make(chan bool, 1)

	go func() {
		var curLog, priorLog *os.File
		for {
			// Open and begin using a new log file
			curLog, _ = os.OpenFile(logPath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0644)
			logger.SetOutput(curLog)

			if priorLog == nil {
				// First log, signal the outer goroutine that we're running
				logging <- true
			} else {
				// Follow-on log, close the prior log file
				priorLog.Close()
				priorLog = nil
			}

			// Wait for a SIGHUP
			<-sigs

			// Setup for the next iteration
			priorLog = curLog
		}
	}()

	signal.Notify(sigs, syscall.SIGHUP)
	<-logging
}

func addStackTraceHandler(logger *log.Logger) {
	sigs := make(chan os.Signal, 1)
	go func() {
		for {
			<-sigs
			buf := make([]byte, 1<<16)
			size := runtime.Stack(buf, true)
			logger.Infof("*** goroutine dump ***\n%s", buf[:size])
		}
	}()
	signal.Notify(sigs, syscall.SIGUSR1)
}
<|endoftext|>


<commit_before>package siesta

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

// Registered services keyed by base URI.
var services = map[string]*Service{}

// A Service is a container for routes with a common base URI.
// It also has two middleware chains, named "pre" and "post".
//
// The "pre" chain is run before the main handler. The first
// handler in the "pre" chain is guaranteed to run, but execution
// may quit anywhere else in the chain.
//
// If the "pre" chain executes completely, the main handler is executed.
// It is skipped otherwise.
//
// The "post" chain runs after the main handler, whether it is skipped
// or not. The first handler in the "post" chain is guaranteed to run, but
// execution may quit anywhere else in the chain if the quit function
// is called.
type Service struct {
	baseURI   string
	trimSlash bool

	pre  []contextHandler
	post []contextHandler

	routes map[string]*node

	notFound contextHandler
}

// NewService returns a new Service with the given base URI
// or panics if the base URI has already been registered.
func NewService(baseURI string) *Service {
	if services[baseURI] != nil {
		panic("service already registered")
	}

	return &Service{
		baseURI:   path.Join("/", baseURI, "/"),
		routes:    map[string]*node{},
		trimSlash: true,
	}
}

// DisableTrimSlash disables the removal of trailing slashes
// before route matching.
func (s *Service) DisableTrimSlash() {
	s.trimSlash = false
}

func addToChain(f interface{}, chain []contextHandler) []contextHandler {
	m := toContextHandler(f)
	return append(chain, m)
}

// AddPre adds f to the end of the "pre" chain.
// It panics if f cannot be converted to a contextHandler (see Service.Route).
func (s *Service) AddPre(f interface{}) {
	s.pre = addToChain(f, s.pre)
}

// AddPost adds f to the end of the "post" chain.
// It panics if f cannot be converted to a contextHandler (see Service.Route).
func (s *Service) AddPost(f interface{}) {
	s.post = addToChain(f, s.post)
}

// Service satisfies the http.Handler interface.
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.ServeHTTPInContext(NewSiestaContext(), w, r)
}

// ServiceHTTPInContext serves an HTTP request within the Context c.
// A Service will run through both of its internal chains, quitting
// when requested.
func (s *Service) ServeHTTPInContext(c Context, w http.ResponseWriter, r *http.Request) {
	r.ParseForm()

	quit := false
	for _, m := range s.pre {
		m(c, w, r, func() {
			quit = true
		})

		if quit {
			// Break out of the "pre" loop, but
			// continue on.
			break
		}
	}

	if !quit {
		// The main handler is only run if we have not
		// been signaled to quit.

		if r.URL.Path != "/" && s.trimSlash {
			r.URL.Path = strings.TrimRight(r.URL.Path, "/")
		}

		var (
			handler contextHandler
			usage   string
			params  routeParams
		)

		// Lookup the tree for this method
		routeNode, ok := s.routes[r.Method]

		if ok {
			handler, usage, params, _ = routeNode.getValue(r.URL.Path)
			c.Set(UsageContextKey, usage)
		}

		if handler == nil {
			if s.notFound != nil {
				// Use user-defined handler.
				s.notFound(c, w, r, func() {})
			} else {
				// Default to the net/http NotFoundHandler.
				http.NotFoundHandler().ServeHTTP(w, r)
			}
		} else {
			for _, p := range params {
				r.Form.Set(p.Key, p.Value)
			}

			handler(c, w, r, func() {
				quit = true
			})
		}
	}

	quit = false
	for _, m := range s.post {
		m(c, w, r, func() {
			quit = true
		})

		if quit {
			return
		}
	}
}

// Route adds a new route to the Service.
// f must be a function with one of the following signatures:
//
//     func(http.ResponseWriter, *http.Request)
//     func(http.ResponseWriter, *http.Request, func())
//     func(Context, http.ResponseWriter, *http.Request)
//     func(Context, http.ResponseWriter, *http.Request, func())
//
// Note that Context is an interface type defined in this package.
// The last argument is a function which is called to signal the
// quitting of the current execution sequence.
func (s *Service) Route(verb, uriPath, usage string, f interface{}) {
	handler := toContextHandler(f)

	if n := s.routes[verb]; n == nil {
		s.routes[verb] = &node{}
	}

	s.routes[verb].addRoute(
		path.Join(s.baseURI, strings.TrimRight(uriPath, "/")),
		usage, handler)
}

// SetNotFound sets the handler for all paths that do not
// match any existing routes. It accepts the same function
// signatures that Route does with the addition of `nil`.
func (s *Service) SetNotFound(f interface{}) {
	if f == nil {
		s.notFound = nil
		return
	}

	handler := toContextHandler(f)
	s.notFound = handler
}

// Register registers s by adding it as a handler to the
// DefaultServeMux in the net/http package.
func (s *Service) Register() {
	http.Handle(s.baseURI, s)
}
<commit_msg>always consume request body<commit_after>package siesta

import (
	"io"
	"io/ioutil"
	"net/http"
	"path"
	"strings"
)

// Registered services keyed by base URI.
var services = map[string]*Service{}

// A Service is a container for routes with a common base URI.
// It also has two middleware chains, named "pre" and "post".
//
// The "pre" chain is run before the main handler. The first
// handler in the "pre" chain is guaranteed to run, but execution
// may quit anywhere else in the chain.
//
// If the "pre" chain executes completely, the main handler is executed.
// It is skipped otherwise.
//
// The "post" chain runs after the main handler, whether it is skipped
// or not. The first handler in the "post" chain is guaranteed to run, but
// execution may quit anywhere else in the chain if the quit function
// is called.
type Service struct {
	baseURI   string
	trimSlash bool

	pre  []contextHandler
	post []contextHandler

	routes map[string]*node

	notFound contextHandler
}

// NewService returns a new Service with the given base URI
// or panics if the base URI has already been registered.
func NewService(baseURI string) *Service {
	if services[baseURI] != nil {
		panic("service already registered")
	}

	return &Service{
		baseURI:   path.Join("/", baseURI, "/"),
		routes:    map[string]*node{},
		trimSlash: true,
	}
}

// DisableTrimSlash disables the removal of trailing slashes
// before route matching.
func (s *Service) DisableTrimSlash() {
	s.trimSlash = false
}

func addToChain(f interface{}, chain []contextHandler) []contextHandler {
	m := toContextHandler(f)
	return append(chain, m)
}

// AddPre adds f to the end of the "pre" chain.
// It panics if f cannot be converted to a contextHandler (see Service.Route).
func (s *Service) AddPre(f interface{}) {
	s.pre = addToChain(f, s.pre)
}

// AddPost adds f to the end of the "post" chain.
// It panics if f cannot be converted to a contextHandler (see Service.Route).
func (s *Service) AddPost(f interface{}) {
	s.post = addToChain(f, s.post)
}

// Service satisfies the http.Handler interface.
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.ServeHTTPInContext(NewSiestaContext(), w, r)
}

// ServiceHTTPInContext serves an HTTP request within the Context c.
// A Service will run through both of its internal chains, quitting
// when requested.
func (s *Service) ServeHTTPInContext(c Context, w http.ResponseWriter, r *http.Request) {
	r.ParseForm()

	quit := false
	for _, m := range s.pre {
		m(c, w, r, func() {
			quit = true
		})

		if quit {
			// Break out of the "pre" loop, but
			// continue on.
			break
		}
	}

	if !quit {
		// The main handler is only run if we have not
		// been signaled to quit.

		if r.URL.Path != "/" && s.trimSlash {
			r.URL.Path = strings.TrimRight(r.URL.Path, "/")
		}

		var (
			handler contextHandler
			usage   string
			params  routeParams
		)

		// Lookup the tree for this method
		routeNode, ok := s.routes[r.Method]

		if ok {
			handler, usage, params, _ = routeNode.getValue(r.URL.Path)
			c.Set(UsageContextKey, usage)
		}

		if handler == nil {
			if s.notFound != nil {
				// Use user-defined handler.
				s.notFound(c, w, r, func() {})
			} else {
				// Default to the net/http NotFoundHandler.
				http.NotFoundHandler().ServeHTTP(w, r)
			}
		} else {
			for _, p := range params {
				r.Form.Set(p.Key, p.Value)
			}

			handler(c, w, r, func() {
				quit = true
			})

			if r.Body != nil {
				io.Copy(ioutil.Discard, r.Body)
				r.Body.Close()
			}
		}
	}

	quit = false
	for _, m := range s.post {
		m(c, w, r, func() {
			quit = true
		})

		if quit {
			return
		}
	}
}

// Route adds a new route to the Service.
// f must be a function with one of the following signatures:
//
//     func(http.ResponseWriter, *http.Request)
//     func(http.ResponseWriter, *http.Request, func())
//     func(Context, http.ResponseWriter, *http.Request)
//     func(Context, http.ResponseWriter, *http.Request, func())
//
// Note that Context is an interface type defined in this package.
// The last argument is a function which is called to signal the
// quitting of the current execution sequence.
func (s *Service) Route(verb, uriPath, usage string, f interface{}) {
	handler := toContextHandler(f)

	if n := s.routes[verb]; n == nil {
		s.routes[verb] = &node{}
	}

	s.routes[verb].addRoute(
		path.Join(s.baseURI, strings.TrimRight(uriPath, "/")),
		usage, handler)
}

// SetNotFound sets the handler for all paths that do not
// match any existing routes. It accepts the same function
// signatures that Route does with the addition of `nil`.
func (s *Service) SetNotFound(f interface{}) {
	if f == nil {
		s.notFound = nil
		return
	}

	handler := toContextHandler(f)
	s.notFound = handler
}

// Register registers s by adding it as a handler to the
// DefaultServeMux in the net/http package.
func (s *Service) Register() {
	http.Handle(s.baseURI, s)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"

	"github.com/Felamande/filesync/log"
	"github.com/kardianos/osext"
)
import (
	"github.com/Felamande/filesync/syncer"
	"github.com/go-martini/martini"
	svc "github.com/kardianos/service"
)

type Program struct {
	Syncer *syncer.Syncer
	Server *martini.ClassicMartini
	Logger *log.FileLogger
	Config *syncer.SavedConfig
	Folder string
}

func (p *Program) Start(s svc.Service) error {
	go p.run()
	return nil
}

func (p *Program) run() {
	folder, err := osext.ExecutableFolder()
	if err != nil {
		fmt.Println(err)
		return
	}

	config, err := ReadConfig(filepath.Join(folder, "config.json"))
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(config)

	p = &Program{
		Config: config,
		Syncer: syncer.New(),
		Server: martini.Classic(),
		Logger: log.NewFileLogger(filepath.Join(folder, "./.log/service.log")),
		Folder: folder,
	}
	p.Server.Map(p.Syncer)
	p.Server.Map(p.Logger)
	p.Server.Post("/new", NewPair)
	p.Server.Get("/new", HelloNewPair)
	go p.Syncer.Run(*p.Config)
	http.ListenAndServe(p.Config.Port, p.Server)
}

func (p *Program) Stop(s svc.Service) error {
	c := syncer.SavedConfig{}
	for _, pair := range p.Syncer.SyncPairs {

		c.Pairs = append(c.Pairs, syncer.SyncPairConfig{
			Left:   pair.Left.Uri(),
			Right:  pair.Right.Uri(),
			Config: pair.Config,
		})
	}
	c.Port = p.Config.Port
	b, err := json.Marshal(&c)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(filepath.Join(p.Folder, "config.json.alt"), b, 0777)
	return err
}

func ReadConfig(ConfigFile string) (*syncer.SavedConfig, error) {

	var config *syncer.SavedConfig = &syncer.SavedConfig{
		Pairs: []syncer.SyncPairConfig{},
		Port:  ":20000",
	}

	configFile, err := os.Open(ConfigFile)
	if err != nil {
		return nil, err
	}
	defer configFile.Close()

	d := json.NewDecoder(configFile)
	err = d.Decode(config)

	if err != nil {
		return nil, err
	}

	if len(config.Port) == 0 {
		config.Port = ":20000"
	}
	if config.Port[0] != ':' {
		config.Port = ":" + config.Port
	}

	return config, nil

}
<commit_msg>Implement system service.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"

	"github.com/Felamande/filesync/log"
	"github.com/kardianos/osext"
)
import (
	"github.com/Felamande/filesync/syncer"
	"github.com/go-martini/martini"
	svc "github.com/kardianos/service"
)

type Program struct {
	Syncer *syncer.Syncer
	Server *martini.ClassicMartini
	Logger *log.FileLogger
	Config *syncer.SavedConfig
	Folder string
}

func (p *Program) Start(s svc.Service) error {
	go p.run()
	return nil
}

func (p *Program) run() {
	folder, err := osext.ExecutableFolder()
	if err != nil {
		fmt.Println(err)
		return
	}

	config, err := ReadConfig(filepath.Join(folder, "config.json"))
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(config)

	p = &Program{
		Config: config,
		Syncer: syncer.New(),
		Server: martini.Classic(),
		Logger: log.NewFileLogger(filepath.Join(folder, "./.log/service.log")),
		Folder: folder,
	}
	p.Server.Map(p.Syncer)
	p.Server.Map(p.Logger)
	p.Server.Post("/new", NewPair)
	p.Server.Get("/new", HelloNewPair)
	go p.Syncer.Run(*p.Config)
	http.ListenAndServe(p.Config.Port, p.Server)
}

func (p *Program) Stop(s svc.Service) error {
	c := syncer.SavedConfig{}
	for _, pair := range p.Syncer.SyncPairs {

		c.Pairs = append(c.Pairs, syncer.SyncPairConfig{
			Left:   pair.Left.Uri(),
			Right:  pair.Right.Uri(),
			Config: pair.Config,
		})
	}
	c.Port = p.Config.Port
	b, err := json.Marshal(&c)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(filepath.Join(p.Folder, "config.json"), b, 0777)
	return err
}

func ReadConfig(ConfigFile string) (*syncer.SavedConfig, error) {

	var config *syncer.SavedConfig = &syncer.SavedConfig{
		Pairs: []syncer.SyncPairConfig{},
		Port:  ":20000",
	}

	configFile, err := os.Open(ConfigFile)
	if err != nil {
		return nil, err
	}
	defer configFile.Close()

	d := json.NewDecoder(configFile)
	err = d.Decode(config)

	if err != nil {
		return nil, err
	}

	if len(config.Port) == 0 {
		config.Port = ":20000"
	}
	if config.Port[0] != ':' {
		config.Port = ":" + config.Port
	}

	return config, nil

}
<|endoftext|>


<commit_before>package micro

import (
	"fmt"
	"os"
	"os/signal"
	rtime "runtime"
	"strings"
	"sync"

	"github.com/micro/go-micro/v2/auth"
	"github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/config/cmd"
	"github.com/micro/go-micro/v2/debug/service/handler"
	"github.com/micro/go-micro/v2/debug/stats"
	"github.com/micro/go-micro/v2/debug/trace"
	"github.com/micro/go-micro/v2/logger"
	"github.com/micro/go-micro/v2/plugin"
	registrySrv "github.com/micro/go-micro/v2/registry/service"
	"github.com/micro/go-micro/v2/runtime"
	"github.com/micro/go-micro/v2/server"
	"github.com/micro/go-micro/v2/store"
	signalutil "github.com/micro/go-micro/v2/util/signal"
	"github.com/micro/go-micro/v2/util/wrapper"
)

type service struct {
	opts Options

	once sync.Once
}

func newService(opts ...Option) Service {
	service := new(service)
	options := newOptions(opts...)

	// service name
	serviceName := options.Server.Options().Name

	// authFn returns the auth, we pass as a function since auth
	// has not yet been set at this point.
	authFn := func() auth.Auth { return options.Server.Options().Auth }

	// wrap client to inject From-Service header on any calls
	options.Client = wrapper.FromService(serviceName, options.Client)
	options.Client = wrapper.TraceCall(serviceName, trace.DefaultTracer, options.Client)
	options.Client = wrapper.AuthClient(serviceName, options.Server.Options().Id, authFn, options.Client)

	// wrap the server to provide handler stats
	options.Server.Init(
		server.WrapHandler(wrapper.HandlerStats(stats.DefaultStats)),
		server.WrapHandler(wrapper.TraceHandler(trace.DefaultTracer)),
		server.WrapHandler(wrapper.AuthHandler(authFn)),
	)

	// set opts
	service.opts = options

	return service
}

func (s *service) Name() string {
	return s.opts.Server.Options().Name
}

// Init initialises options. Additionally it calls cmd.Init
// which parses command line flags. cmd.Init is only called
// on first Init.
func (s *service) Init(opts ...Option) {
	// process options
	for _, o := range opts {
		o(&s.opts)
	}

	s.once.Do(func() {
		// setup the plugins
		for _, p := range strings.Split(os.Getenv("MICRO_PLUGIN"), ",") {
			if len(p) == 0 {
				continue
			}

			// load the plugin
			c, err := plugin.Load(p)
			if err != nil {
				logger.Fatal(err)
			}

			// initialise the plugin
			if err := plugin.Init(c); err != nil {
				logger.Fatal(err)
			}
		}

		// set cmd name
		if len(s.opts.Cmd.App().Name) == 0 {
			s.opts.Cmd.App().Name = s.Server().Options().Name
		}

		// Initialise the command flags, overriding new service
		if err := s.opts.Cmd.Init(
			cmd.Auth(&s.opts.Auth),
			cmd.Broker(&s.opts.Broker),
			cmd.Registry(&s.opts.Registry),
			cmd.Runtime(&s.opts.Runtime),
			cmd.Transport(&s.opts.Transport),
			cmd.Client(&s.opts.Client),
			cmd.Config(&s.opts.Config),
			cmd.Server(&s.opts.Server),
			cmd.Store(&s.opts.Store),
			cmd.Profile(&s.opts.Profile),
		); err != nil {
			logger.Fatal(err)
		}

		// Explicitly set the table name to the service name
		name := s.opts.Cmd.App().Name
		s.opts.Store.Init(store.Table(name))

		// Set the client for the micro clients
		s.opts.Auth.Init(auth.WithClient(s.Client()))
		s.opts.Registry.Init(registrySrv.WithClient(s.Client()))
		s.opts.Runtime.Init(runtime.WithClient(s.Client()))
		s.opts.Store.Init(store.WithClient(s.Client()))
	})
}

func (s *service) Options() Options {
	return s.opts
}

func (s *service) Client() client.Client {
	return s.opts.Client
}

func (s *service) Server() server.Server {
	return s.opts.Server
}

func (s *service) String() string {
	return "micro"
}

func (s *service) Start() error {
	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	if err := s.opts.Server.Start(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	return nil
}

func (s *service) Stop() error {
	var gerr error

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	if err := s.opts.Server.Stop(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	return gerr
}

func (s *service) Run() error {
	// register the debug handler
	s.opts.Server.Handle(
		s.opts.Server.NewHandler(
			handler.NewHandler(),
			server.InternalHandler(true),
		),
	)

	// start the profiler
	if s.opts.Profile != nil {
		// to view mutex contention
		rtime.SetMutexProfileFraction(5)
		// to view blocking profile
		rtime.SetBlockProfileRate(1)

		if err := s.opts.Profile.Start(); err != nil {
			return err
		}
		defer s.opts.Profile.Stop()
	}

	if logger.V(logger.InfoLevel, logger.DefaultLogger) {
		logger.Infof("Starting [service] %s", s.Name())
	}

	// generate an auth account
	if err := s.registerAuthAccount(); err != nil {
		return err
	}

	if err := s.Start(); err != nil {
		return err
	}

	ch := make(chan os.Signal, 1)
	if s.opts.Signal {
		signal.Notify(ch, signalutil.Shutdown()...)
	}

	select {
	// wait on kill signal
	case <-ch:
	// wait on context cancel
	case <-s.opts.Context.Done():
	}

	return s.Stop()
}

func (s *service) registerAuthAccount() error {
	// determine the type of service from the name. we do this so we can allocate
	// different roles depending on the type of services. e.g. we don't want web
	// services talking directly to the runtime. TODO: find a better way to determine
	// the type of service
	serviceType := "service"
	if strings.Contains(s.Name(), "api") {
		serviceType = "api"
	} else if strings.Contains(s.Name(), "web") {
		serviceType = "web"
	}

	// generate a new auth account for the service
	name := fmt.Sprintf("%v-%v", s.Name(), s.Server().Options().Id)
	opts := []auth.GenerateOption{
		auth.WithRoles(serviceType),
		auth.WithNamespace(s.Options().Auth.Options().Namespace),
	}
	acc, err := s.Options().Auth.Generate(name, opts...)
	if err != nil {
		return err
	}

	// generate a token
	token, err := s.Options().Auth.Token(auth.WithCredentials(acc.ID, acc.Secret))
	if err != nil {
		return err
	}
	s.Options().Auth.Init(auth.ClientToken(token))

	logger.Infof("Auth [%v] Authenticated as %v", s.Options().Auth, name)
	return nil
}
<commit_msg>Remove service type role<commit_after>package micro

import (
	"fmt"
	"os"
	"os/signal"
	rtime "runtime"
	"strings"
	"sync"

	"github.com/micro/go-micro/v2/auth"
	"github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/config/cmd"
	"github.com/micro/go-micro/v2/debug/service/handler"
	"github.com/micro/go-micro/v2/debug/stats"
	"github.com/micro/go-micro/v2/debug/trace"
	"github.com/micro/go-micro/v2/logger"
	"github.com/micro/go-micro/v2/plugin"
	registrySrv "github.com/micro/go-micro/v2/registry/service"
	"github.com/micro/go-micro/v2/runtime"
	"github.com/micro/go-micro/v2/server"
	"github.com/micro/go-micro/v2/store"
	signalutil "github.com/micro/go-micro/v2/util/signal"
	"github.com/micro/go-micro/v2/util/wrapper"
)

type service struct {
	opts Options

	once sync.Once
}

func newService(opts ...Option) Service {
	service := new(service)
	options := newOptions(opts...)

	// service name
	serviceName := options.Server.Options().Name

	// authFn returns the auth, we pass as a function since auth
	// has not yet been set at this point.
	authFn := func() auth.Auth { return options.Server.Options().Auth }

	// wrap client to inject From-Service header on any calls
	options.Client = wrapper.FromService(serviceName, options.Client)
	options.Client = wrapper.TraceCall(serviceName, trace.DefaultTracer, options.Client)
	options.Client = wrapper.AuthClient(serviceName, options.Server.Options().Id, authFn, options.Client)

	// wrap the server to provide handler stats
	options.Server.Init(
		server.WrapHandler(wrapper.HandlerStats(stats.DefaultStats)),
		server.WrapHandler(wrapper.TraceHandler(trace.DefaultTracer)),
		server.WrapHandler(wrapper.AuthHandler(authFn)),
	)

	// set opts
	service.opts = options

	return service
}

func (s *service) Name() string {
	return s.opts.Server.Options().Name
}

// Init initialises options. Additionally it calls cmd.Init
// which parses command line flags. cmd.Init is only called
// on first Init.
func (s *service) Init(opts ...Option) {
	// process options
	for _, o := range opts {
		o(&s.opts)
	}

	s.once.Do(func() {
		// setup the plugins
		for _, p := range strings.Split(os.Getenv("MICRO_PLUGIN"), ",") {
			if len(p) == 0 {
				continue
			}

			// load the plugin
			c, err := plugin.Load(p)
			if err != nil {
				logger.Fatal(err)
			}

			// initialise the plugin
			if err := plugin.Init(c); err != nil {
				logger.Fatal(err)
			}
		}

		// set cmd name
		if len(s.opts.Cmd.App().Name) == 0 {
			s.opts.Cmd.App().Name = s.Server().Options().Name
		}

		// Initialise the command flags, overriding new service
		if err := s.opts.Cmd.Init(
			cmd.Auth(&s.opts.Auth),
			cmd.Broker(&s.opts.Broker),
			cmd.Registry(&s.opts.Registry),
			cmd.Runtime(&s.opts.Runtime),
			cmd.Transport(&s.opts.Transport),
			cmd.Client(&s.opts.Client),
			cmd.Config(&s.opts.Config),
			cmd.Server(&s.opts.Server),
			cmd.Store(&s.opts.Store),
			cmd.Profile(&s.opts.Profile),
		); err != nil {
			logger.Fatal(err)
		}

		// Explicitly set the table name to the service name
		name := s.opts.Cmd.App().Name
		s.opts.Store.Init(store.Table(name))

		// Set the client for the micro clients
		s.opts.Auth.Init(auth.WithClient(s.Client()))
		s.opts.Registry.Init(registrySrv.WithClient(s.Client()))
		s.opts.Runtime.Init(runtime.WithClient(s.Client()))
		s.opts.Store.Init(store.WithClient(s.Client()))
	})
}

func (s *service) Options() Options {
	return s.opts
}

func (s *service) Client() client.Client {
	return s.opts.Client
}

func (s *service) Server() server.Server {
	return s.opts.Server
}

func (s *service) String() string {
	return "micro"
}

func (s *service) Start() error {
	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	if err := s.opts.Server.Start(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	return nil
}

func (s *service) Stop() error {
	var gerr error

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	if err := s.opts.Server.Stop(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	return gerr
}

func (s *service) Run() error {
	// register the debug handler
	s.opts.Server.Handle(
		s.opts.Server.NewHandler(
			handler.NewHandler(),
			server.InternalHandler(true),
		),
	)

	// start the profiler
	if s.opts.Profile != nil {
		// to view mutex contention
		rtime.SetMutexProfileFraction(5)
		// to view blocking profile
		rtime.SetBlockProfileRate(1)

		if err := s.opts.Profile.Start(); err != nil {
			return err
		}
		defer s.opts.Profile.Stop()
	}

	if logger.V(logger.InfoLevel, logger.DefaultLogger) {
		logger.Infof("Starting [service] %s", s.Name())
	}

	// generate an auth account
	if err := s.registerAuthAccount(); err != nil {
		return err
	}

	if err := s.Start(); err != nil {
		return err
	}

	ch := make(chan os.Signal, 1)
	if s.opts.Signal {
		signal.Notify(ch, signalutil.Shutdown()...)
	}

	select {
	// wait on kill signal
	case <-ch:
	// wait on context cancel
	case <-s.opts.Context.Done():
	}

	return s.Stop()
}

func (s *service) registerAuthAccount() error {
	// generate a new auth account for the service
	name := fmt.Sprintf("%v-%v", s.Name(), s.Server().Options().Id)
	opts := []auth.GenerateOption{
		auth.WithType("service"),
		auth.WithRoles("service"),
		auth.WithNamespace(s.Options().Auth.Options().Namespace),
	}
	acc, err := s.Options().Auth.Generate(name, opts...)
	if err != nil {
		return err
	}

	// generate a token
	token, err := s.Options().Auth.Token(auth.WithCredentials(acc.ID, acc.Secret))
	if err != nil {
		return err
	}
	s.Options().Auth.Init(auth.ClientToken(token))

	logger.Infof("Auth [%v] Authenticated as %v", s.Options().Auth, name)
	return nil
}
<|endoftext|>


<commit_before>// session.go - mixnet client session
// Copyright (C) 2018  David Stainton.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package client

import (
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	mrand "math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/katzenpost/client/config"
	cConstants "github.com/katzenpost/client/constants"
	"github.com/katzenpost/client/internal/pkiclient"
	"github.com/katzenpost/client/utils"
	coreConstants "github.com/katzenpost/core/constants"
	"github.com/katzenpost/core/crypto/ecdh"
	"github.com/katzenpost/core/log"
	"github.com/katzenpost/core/pki"
	"github.com/katzenpost/core/sphinx"
	sConstants "github.com/katzenpost/core/sphinx/constants"
	"github.com/katzenpost/core/worker"
	"github.com/katzenpost/minclient"
	"gopkg.in/eapache/channels.v1"
	"gopkg.in/op/go-logging.v1"
)

// Session is the struct type that keeps state for a given session.
type Session struct {
	worker.Worker

	cfg       *config.Config
	pkiClient pki.Client
	minclient *minclient.Client
	log       *logging.Logger

	fatalErrCh chan error
	opCh       chan workerOp

	eventCh   channels.Channel
	EventSink chan Event

	linkKey   *ecdh.PrivateKey
	onlineAt  time.Time
	hasPKIDoc bool

	egressQueue EgressQueue
	rescheduler *rescheduler

	surbIDMap        sync.Map // [sConstants.SURBIDLength]byte -> *Message
	sentWaitChanMap  sync.Map // MessageID -> chan *Message
	replyWaitChanMap sync.Map // MessageID -> chan []byte

	decoyLoopTally uint64
}

// New establishes a session with provider using key.
// This method will block until session is connected to the Provider.
func NewSession(
	ctx context.Context,
	fatalErrCh chan error,
	logBackend *log.Backend,
	cfg *config.Config,
	linkKey *ecdh.PrivateKey) (*Session, error) {
	var err error

	// create a pkiclient for our own client lookups
	// AND create a pkiclient for minclient's use
	proxyCfg := cfg.UpstreamProxyConfig()
	pkiClient, err := cfg.NewPKIClient(logBackend, proxyCfg)
	if err != nil {
		return nil, err
	}

	// create a pkiclient for minclient's use
	pkiClient2, err := cfg.NewPKIClient(logBackend, proxyCfg)
	if err != nil {
		return nil, err
	}
	pkiCacheClient := pkiclient.New(pkiClient2)

	clientLog := logBackend.GetLogger(fmt.Sprintf("%s@%s_client", cfg.Account.User, cfg.Account.Provider))

	s := &Session{
		cfg:         cfg,
		linkKey:     linkKey,
		pkiClient:   pkiClient,
		log:         clientLog,
		fatalErrCh:  fatalErrCh,
		eventCh:     channels.NewInfiniteChannel(),
		EventSink:   make(chan Event),
		opCh:        make(chan workerOp, 8),
		egressQueue: new(Queue),
	}
	// Configure the rescheduler instance
	s.rescheduler = NewRescheduler(s)
	// Configure and bring up the minclient instance.
	clientCfg := &minclient.ClientConfig{
		User:                cfg.Account.User,
		Provider:            cfg.Account.Provider,
		ProviderKeyPin:      cfg.Account.ProviderKeyPin,
		LinkKey:             s.linkKey,
		LogBackend:          logBackend,
		PKIClient:           pkiCacheClient,
		OnConnFn:            s.onConnection,
		OnMessageFn:         s.onMessage,
		OnACKFn:             s.onACK,
		OnDocumentFn:        s.onDocument,
		DialContextFn:       proxyCfg.ToDialContext("authority"),
		PreferedTransports:  cfg.Debug.PreferedTransports,
		MessagePollInterval: time.Duration(cfg.Debug.PollingInterval) * time.Millisecond,
		EnableTimeSync:      false, // Be explicit about it.
	}

	s.Go(s.eventSinkWorker)
	s.Go(s.garbageCollectionWorker)

	s.minclient, err = minclient.New(clientCfg)
	if err != nil {
		return nil, err
	}

	// block until we get the first PKI document
	// and then set our timers accordingly
	err = s.awaitFirstPKIDoc(ctx)
	if err != nil {
		return nil, err
	}
	s.Go(s.worker)
	return s, nil
}

func (s *Session) eventSinkWorker() {
	for {
		select {
		case <-s.HaltCh():
			s.log.Debugf("Event sink worker terminating gracefully.")
			return
		case e := <-s.eventCh.Out():
			select {
			case s.EventSink <- e.(Event):
			case <-s.HaltCh():
				s.log.Debugf("Event sink worker terminating gracefully.")
				return
			}
		}
	}
}

func (s *Session) garbageCollectionWorker() {
	timer := time.NewTimer(cConstants.GarbageCollectionInterval)
	defer timer.Stop()
	for {
		select {
		case <-s.HaltCh():
			s.log.Debugf("Garbage collection worker terminating gracefully.")
			return
		case <-timer.C:
			s.garbageCollect()
			timer.Reset(cConstants.GarbageCollectionInterval)
		}
	}
}

func (s *Session) garbageCollect() {
	s.log.Debug("Running garbage collection process.")
	// [sConstants.SURBIDLength]byte -> *Message
	surbIDMapRange := func(rawSurbID, rawMessage interface{}) bool {
		surbID := rawSurbID.([sConstants.SURBIDLength]byte)
		message := rawMessage.(*Message)
		if time.Now().After(message.SentAt.Add(message.ReplyETA).Add(cConstants.RoundTripTimeSlop)) {
			s.log.Debug("Garbage collecting SURB ID Map entry for Message ID %x", message.ID)
			s.surbIDMap.Delete(surbID)
			s.eventCh.In() <- &MessageIDGarbageCollected{
				MessageID: message.ID,
			}
		}
		return true
	}
	s.surbIDMap.Range(surbIDMapRange)
}

func (s *Session) awaitFirstPKIDoc(ctx context.Context) error {
	for {
		var qo workerOp
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-s.HaltCh():
			s.log.Debugf("Await first pki doc worker terminating gracefully")
			return errors.New("terminating gracefully")
		case <-time.After(time.Duration(s.cfg.Debug.InitialMaxPKIRetrievalDelay) * time.Second):
			return errors.New("timeout failure awaiting first PKI document")
		case qo = <-s.opCh:
		}
		switch op := qo.(type) {
		case opNewDocument:
			// Determine if PKI doc is valid. If not then abort.
			err := s.isDocValid(op.doc)
			if err != nil {
				s.fatalErrCh <- fmt.Errorf("aborting, PKI doc is not valid for our decoy traffic use case: %v", err)
				return err
			}
			s.setPollIntervalFromDoc(op.doc)
			return nil
		default:
			continue
		}
	}
	// NOT REACHED
}

// GetService returns a randomly selected service
// matching the specified service name
func (s *Session) GetService(serviceName string) (*utils.ServiceDescriptor, error) {
	doc := s.minclient.CurrentDocument()
	if doc == nil {
		return nil, errors.New("pki doc is nil")
	}
	serviceDescriptors := utils.FindServices(serviceName, doc)
	if len(serviceDescriptors) == 0 {
		return nil, errors.New("error, GetService failure, service not found in pki doc")
	}
	return &serviceDescriptors[mrand.Intn(len(serviceDescriptors))], nil
}

// OnConnection will be called by the minclient api
// upon connection change status to the Provider
func (s *Session) onConnection(err error) {
	s.log.Debugf("onConnection %v", err)
	s.eventCh.In() <- &ConnectionStatusEvent{
		IsConnected: err == nil,
		Err:         err,
	}
	s.opCh <- opConnStatusChanged{
		isConnected: err == nil,
	}
}

// OnMessage will be called by the minclient api
// upon receiving a message
func (s *Session) onMessage(ciphertextBlock []byte) error {
	s.log.Debugf("OnMessage")
	return nil
}

func (s *Session) incrementDecoyLoopTally() {
	atomic.AddUint64(&s.decoyLoopTally, 1)
}

func (s *Session) decrementDecoyLoopTally() {
	atomic.AddUint64(&s.decoyLoopTally, ^uint64(0))
}

// OnACK is called by the minclient api when we receive a SURB reply message.
func (s *Session) onACK(surbID *[sConstants.SURBIDLength]byte, ciphertext []byte) error {
	idStr := fmt.Sprintf("[%v]", hex.EncodeToString(surbID[:]))
	s.log.Infof("OnACK with SURBID %x", idStr)

	rawMessage, ok := s.surbIDMap.Load(*surbID)
	if !ok {
		s.log.Debug("Strange, received reply with unexpected SURBID")
		return nil
	}
	s.surbIDMap.Delete(*surbID)
	msg := rawMessage.(*Message)
	plaintext, err := sphinx.DecryptSURBPayload(ciphertext, msg.Key)
	if err != nil {
		s.log.Infof("Discarding SURB Reply, decryption failure: %s", err)
		return nil
	}
	if len(plaintext) != coreConstants.ForwardPayloadLength {
		s.log.Warningf("Discarding SURB %v: Invalid payload size: %v", idStr, len(plaintext))
		return nil
	}
	if msg.WithSURB && msg.IsDecoy {
		s.decrementDecoyLoopTally()
		return nil
	}

	if msg.IsBlocking {
		replyWaitChanRaw, ok := s.replyWaitChanMap.Load(*msg.ID)
		if !ok {
			//XXX: this can happen if a SURB-ACK arrives after a call to BlockingSendUnreliableMessage has timed-out
			// because the session.surbIDMap has not been deleted or garbage collected
			s.log.Warningf("Discarding surb %v for blocking message %x : caller likely timed-out", idStr, msg.ID)
			return nil
		}
		replyWaitChan := replyWaitChanRaw.(chan []byte)
		replyWaitChan <- plaintext[2:]
	} else {
		s.eventCh.In() <- &MessageReplyEvent{
			MessageID: msg.ID,
			Payload:   plaintext[2:],
			Err:       nil,
		}
	}
	return nil
}

func (s *Session) onDocument(doc *pki.Document) {
	s.log.Debugf("onDocument(): Epoch %v", doc.Epoch)
	s.hasPKIDoc = true
	s.opCh <- opNewDocument{
		doc: doc,
	}
	s.eventCh.In() <- &NewDocumentEvent{
		Document: doc,
	}
}

func (s *Session) CurrentDocument() *pki.Document {
	return s.minclient.CurrentDocument()
}

func (s *Session) GetReunionConfig() *config.Reunion {
	return s.cfg.Reunion
}

func (s *Session) GetPandaConfig() *config.Panda {
	return s.cfg.Panda
}

func (s *Session) Shutdown() {
	s.Halt()
	s.minclient.Shutdown()
	s.minclient.Wait()
}
<commit_msg>fix OnACK log<commit_after>// session.go - mixnet client session
// Copyright (C) 2018  David Stainton.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package client

import (
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	mrand "math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/katzenpost/client/config"
	cConstants "github.com/katzenpost/client/constants"
	"github.com/katzenpost/client/internal/pkiclient"
	"github.com/katzenpost/client/utils"
	coreConstants "github.com/katzenpost/core/constants"
	"github.com/katzenpost/core/crypto/ecdh"
	"github.com/katzenpost/core/log"
	"github.com/katzenpost/core/pki"
	"github.com/katzenpost/core/sphinx"
	sConstants "github.com/katzenpost/core/sphinx/constants"
	"github.com/katzenpost/core/worker"
	"github.com/katzenpost/minclient"
	"gopkg.in/eapache/channels.v1"
	"gopkg.in/op/go-logging.v1"
)

// Session is the struct type that keeps state for a given session.
type Session struct {
	worker.Worker

	cfg       *config.Config
	pkiClient pki.Client
	minclient *minclient.Client
	log       *logging.Logger

	fatalErrCh chan error
	opCh       chan workerOp

	eventCh   channels.Channel
	EventSink chan Event

	linkKey   *ecdh.PrivateKey
	onlineAt  time.Time
	hasPKIDoc bool

	egressQueue EgressQueue
	rescheduler *rescheduler

	surbIDMap        sync.Map // [sConstants.SURBIDLength]byte -> *Message
	sentWaitChanMap  sync.Map // MessageID -> chan *Message
	replyWaitChanMap sync.Map // MessageID -> chan []byte

	decoyLoopTally uint64
}

// New establishes a session with provider using key.
// This method will block until session is connected to the Provider.
func NewSession(
	ctx context.Context,
	fatalErrCh chan error,
	logBackend *log.Backend,
	cfg *config.Config,
	linkKey *ecdh.PrivateKey) (*Session, error) {
	var err error

	// create a pkiclient for our own client lookups
	// AND create a pkiclient for minclient's use
	proxyCfg := cfg.UpstreamProxyConfig()
	pkiClient, err := cfg.NewPKIClient(logBackend, proxyCfg)
	if err != nil {
		return nil, err
	}

	// create a pkiclient for minclient's use
	pkiClient2, err := cfg.NewPKIClient(logBackend, proxyCfg)
	if err != nil {
		return nil, err
	}
	pkiCacheClient := pkiclient.New(pkiClient2)

	clientLog := logBackend.GetLogger(fmt.Sprintf("%s@%s_client", cfg.Account.User, cfg.Account.Provider))

	s := &Session{
		cfg:         cfg,
		linkKey:     linkKey,
		pkiClient:   pkiClient,
		log:         clientLog,
		fatalErrCh:  fatalErrCh,
		eventCh:     channels.NewInfiniteChannel(),
		EventSink:   make(chan Event),
		opCh:        make(chan workerOp, 8),
		egressQueue: new(Queue),
	}
	// Configure the rescheduler instance
	s.rescheduler = NewRescheduler(s)
	// Configure and bring up the minclient instance.
	clientCfg := &minclient.ClientConfig{
		User:                cfg.Account.User,
		Provider:            cfg.Account.Provider,
		ProviderKeyPin:      cfg.Account.ProviderKeyPin,
		LinkKey:             s.linkKey,
		LogBackend:          logBackend,
		PKIClient:           pkiCacheClient,
		OnConnFn:            s.onConnection,
		OnMessageFn:         s.onMessage,
		OnACKFn:             s.onACK,
		OnDocumentFn:        s.onDocument,
		DialContextFn:       proxyCfg.ToDialContext("authority"),
		PreferedTransports:  cfg.Debug.PreferedTransports,
		MessagePollInterval: time.Duration(cfg.Debug.PollingInterval) * time.Millisecond,
		EnableTimeSync:      false, // Be explicit about it.
	}

	s.Go(s.eventSinkWorker)
	s.Go(s.garbageCollectionWorker)

	s.minclient, err = minclient.New(clientCfg)
	if err != nil {
		return nil, err
	}

	// block until we get the first PKI document
	// and then set our timers accordingly
	err = s.awaitFirstPKIDoc(ctx)
	if err != nil {
		return nil, err
	}
	s.Go(s.worker)
	return s, nil
}

func (s *Session) eventSinkWorker() {
	for {
		select {
		case <-s.HaltCh():
			s.log.Debugf("Event sink worker terminating gracefully.")
			return
		case e := <-s.eventCh.Out():
			select {
			case s.EventSink <- e.(Event):
			case <-s.HaltCh():
				s.log.Debugf("Event sink worker terminating gracefully.")
				return
			}
		}
	}
}

func (s *Session) garbageCollectionWorker() {
	timer := time.NewTimer(cConstants.GarbageCollectionInterval)
	defer timer.Stop()
	for {
		select {
		case <-s.HaltCh():
			s.log.Debugf("Garbage collection worker terminating gracefully.")
			return
		case <-timer.C:
			s.garbageCollect()
			timer.Reset(cConstants.GarbageCollectionInterval)
		}
	}
}

func (s *Session) garbageCollect() {
	s.log.Debug("Running garbage collection process.")
	// [sConstants.SURBIDLength]byte -> *Message
	surbIDMapRange := func(rawSurbID, rawMessage interface{}) bool {
		surbID := rawSurbID.([sConstants.SURBIDLength]byte)
		message := rawMessage.(*Message)
		if time.Now().After(message.SentAt.Add(message.ReplyETA).Add(cConstants.RoundTripTimeSlop)) {
			s.log.Debug("Garbage collecting SURB ID Map entry for Message ID %x", message.ID)
			s.surbIDMap.Delete(surbID)
			s.eventCh.In() <- &MessageIDGarbageCollected{
				MessageID: message.ID,
			}
		}
		return true
	}
	s.surbIDMap.Range(surbIDMapRange)
}

func (s *Session) awaitFirstPKIDoc(ctx context.Context) error {
	for {
		var qo workerOp
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-s.HaltCh():
			s.log.Debugf("Await first pki doc worker terminating gracefully")
			return errors.New("terminating gracefully")
		case <-time.After(time.Duration(s.cfg.Debug.InitialMaxPKIRetrievalDelay) * time.Second):
			return errors.New("timeout failure awaiting first PKI document")
		case qo = <-s.opCh:
		}
		switch op := qo.(type) {
		case opNewDocument:
			// Determine if PKI doc is valid. If not then abort.
			err := s.isDocValid(op.doc)
			if err != nil {
				s.fatalErrCh <- fmt.Errorf("aborting, PKI doc is not valid for our decoy traffic use case: %v", err)
				return err
			}
			s.setPollIntervalFromDoc(op.doc)
			return nil
		default:
			continue
		}
	}
	// NOT REACHED
}

// GetService returns a randomly selected service
// matching the specified service name
func (s *Session) GetService(serviceName string) (*utils.ServiceDescriptor, error) {
	doc := s.minclient.CurrentDocument()
	if doc == nil {
		return nil, errors.New("pki doc is nil")
	}
	serviceDescriptors := utils.FindServices(serviceName, doc)
	if len(serviceDescriptors) == 0 {
		return nil, errors.New("error, GetService failure, service not found in pki doc")
	}
	return &serviceDescriptors[mrand.Intn(len(serviceDescriptors))], nil
}

// OnConnection will be called by the minclient api
// upon connection change status to the Provider
func (s *Session) onConnection(err error) {
	s.log.Debugf("onConnection %v", err)
	s.eventCh.In() <- &ConnectionStatusEvent{
		IsConnected: err == nil,
		Err:         err,
	}
	s.opCh <- opConnStatusChanged{
		isConnected: err == nil,
	}
}

// OnMessage will be called by the minclient api
// upon receiving a message
func (s *Session) onMessage(ciphertextBlock []byte) error {
	s.log.Debugf("OnMessage")
	return nil
}

func (s *Session) incrementDecoyLoopTally() {
	atomic.AddUint64(&s.decoyLoopTally, 1)
}

func (s *Session) decrementDecoyLoopTally() {
	atomic.AddUint64(&s.decoyLoopTally, ^uint64(0))
}

// OnACK is called by the minclient api when we receive a SURB reply message.
func (s *Session) onACK(surbID *[sConstants.SURBIDLength]byte, ciphertext []byte) error {
	idStr := fmt.Sprintf("[%v]", hex.EncodeToString(surbID[:]))
	s.log.Infof("OnACK with SURBID %s", idStr)

	rawMessage, ok := s.surbIDMap.Load(*surbID)
	if !ok {
		s.log.Debug("Strange, received reply with unexpected SURBID")
		return nil
	}
	s.surbIDMap.Delete(*surbID)
	msg := rawMessage.(*Message)
	plaintext, err := sphinx.DecryptSURBPayload(ciphertext, msg.Key)
	if err != nil {
		s.log.Infof("Discarding SURB Reply, decryption failure: %s", err)
		return nil
	}
	if len(plaintext) != coreConstants.ForwardPayloadLength {
		s.log.Warningf("Discarding SURB %v: Invalid payload size: %v", idStr, len(plaintext))
		return nil
	}
	if msg.WithSURB && msg.IsDecoy {
		s.decrementDecoyLoopTally()
		return nil
	}

	if msg.IsBlocking {
		replyWaitChanRaw, ok := s.replyWaitChanMap.Load(*msg.ID)
		if !ok {
			//XXX: this can happen if a SURB-ACK arrives after a call to BlockingSendUnreliableMessage has timed-out
			// because the session.surbIDMap has not been deleted or garbage collected
			s.log.Warningf("Discarding surb %v for blocking message %x : caller likely timed-out", idStr, msg.ID)
			return nil
		}
		replyWaitChan := replyWaitChanRaw.(chan []byte)
		replyWaitChan <- plaintext[2:]
	} else {
		s.eventCh.In() <- &MessageReplyEvent{
			MessageID: msg.ID,
			Payload:   plaintext[2:],
			Err:       nil,
		}
	}
	return nil
}

func (s *Session) onDocument(doc *pki.Document) {
	s.log.Debugf("onDocument(): Epoch %v", doc.Epoch)
	s.hasPKIDoc = true
	s.opCh <- opNewDocument{
		doc: doc,
	}
	s.eventCh.In() <- &NewDocumentEvent{
		Document: doc,
	}
}

func (s *Session) CurrentDocument() *pki.Document {
	return s.minclient.CurrentDocument()
}

func (s *Session) GetReunionConfig() *config.Reunion {
	return s.cfg.Reunion
}

func (s *Session) GetPandaConfig() *config.Panda {
	return s.cfg.Panda
}

func (s *Session) Shutdown() {
	s.Halt()
	s.minclient.Shutdown()
	s.minclient.Wait()
}
<|endoftext|>


<commit_before>package yamux

import (
	"fmt"
	"io"
	"math"
	"net"
	"sync"
	"time"
)

// Session is used to wrap a reliable ordered connection and to
// multiplex it into multiple streams.
type Session struct {
	// client is true if we are a client size connection
	client bool

	// config holds our configuration
	config *Config

	// conn is the underlying connection
	conn io.ReadWriteCloser

	// pings is used to track inflight pings
	pings    map[uint32]chan struct{}
	pingID   uint32
	pingLock sync.Mutex

	// remoteGoAway indicates the remote side does
	// not want futher connections
	remoteGoAway bool

	// localGoAway indicates that we should stop
	// accepting futher connections
	localGoAway bool

	// nextStreamID is the next stream we should
	// send. This depends if we are a client/server.
	nextStreamID uint32

	// streams maps a stream id to a stream
	streams    map[uint32]*Stream
	streamLock sync.RWMutex

	// acceptCh is used to pass ready streams to the client
	acceptCh chan *Stream

	// sendCh is used to mark a stream as ready to send,
	// or to send a header out directly.
	sendCh chan sendReady

	// shutdown is used to safely close a session
	shutdown     bool
	shutdownErr  error
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// sendReady is used to either mark a stream as ready
// or to directly send a header
type sendReady struct {
	Hdr  []byte
	Body io.Reader
	Err  chan error
}

// newSession is used to construct a new session
func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := &Session{
		client:     client,
		config:     config,
		conn:       conn,
		pings:      make(map[uint32]chan struct{}),
		streams:    make(map[uint32]*Stream),
		acceptCh:   make(chan *Stream, config.AcceptBacklog),
		sendCh:     make(chan sendReady, 64),
		shutdownCh: make(chan struct{}),
	}
	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 2
	}
	go s.recv()
	go s.send()
	if config.EnableKeepAlive {
		go s.keepalive()
	}
	return s
}

// isShutdown does a safe check to see if we have shutdown
func (s *Session) isShutdown() bool {
	select {
	case <-s.shutdownCh:
		return true
	default:
		return false
	}
}

// Open is used to create a new stream
func (s *Session) Open() (*Stream, error) {
	if s.isShutdown() {
		return nil, ErrSessionShutdown
	}
	if s.remoteGoAway {
		return nil, ErrRemoteGoAway
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()

	// Check if we've exhaused the streams
	id := s.nextStreamID
	if id >= math.MaxUint32-1 {
		return nil, ErrStreamsExhausted
	}
	s.nextStreamID += 2

	// Register the stream
	stream := newStream(s, id, streamInit)
	s.streams[id] = stream

	// Send the window update to create
	return stream, stream.sendWindowUpdate()
}

// Accept is used to block until the next available stream
// is ready to be accepted.
func (s *Session) Accept() (net.Conn, error) {
	return s.AcceptStream()
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	select {
	case stream := <-s.acceptCh:
		return stream, nil
	case <-s.shutdownCh:
		return nil, s.shutdownErr
	}
}

// Close is used to close the session and all streams.
// Attempts to send a GoAway before closing the connection.
func (s *Session) Close() error {
	s.shutdownLock.Lock()
	defer s.shutdownLock.Unlock()

	if s.shutdown {
		return nil
	}
	s.shutdown = true
	if s.shutdownErr == nil {
		s.shutdownErr = ErrSessionShutdown
	}
	close(s.shutdownCh)
	s.conn.Close()

	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	for _, stream := range s.streams {
		stream.forceClose()
	}
	return nil
}

// GoAway can be used to prevent accepting further
// connections. It does not close the underlying conn.
func (s *Session) GoAway() error {
	s.localGoAway = true
	s.goAway(goAwayNormal)
	return nil
}

// Ping is used to measure the RTT response time
func (s *Session) Ping() (time.Duration, error) {
	// Get a channel for the ping
	ch := make(chan struct{})

	// Get a new ping id, mark as pending
	s.pingLock.Lock()
	id := s.pingID
	s.pingID++
	s.pings[id] = ch
	s.pingLock.Unlock()

	// Send the ping request
	hdr := header(make([]byte, headerSize))
	hdr.encode(typePing, flagSYN, 0, id)
	if err := s.waitForSend(hdr, nil); err != nil {
		return 0, err
	}

	// Wait for a response
	start := time.Now()
	select {
	case <-ch:
	case <-s.shutdownCh:
		return 0, ErrSessionShutdown
	}

	// Compute the RTT
	return time.Now().Sub(start), nil
}

// keepalive is a long running goroutine that periodically does
// a ping to keep the connection alive.
func (s *Session) keepalive() {
	for {
		select {
		case <-time.After(s.config.KeepAliveInterval):
			s.Ping()
		case <-s.shutdownCh:
			return
		}
	}
}

// waitForSend waits to send a header, checking for a potential shutdown
func (s *Session) waitForSend(hdr header, body io.Reader) error {
	errCh := make(chan error, 1)
	ready := sendReady{Hdr: hdr, Body: body, Err: errCh}
	select {
	case s.sendCh <- ready:
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
	select {
	case err := <-errCh:
		return err
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
}

// sendNoWait does a send without waiting
func (s *Session) sendNoWait(hdr header) error {
	select {
	case s.sendCh <- sendReady{Hdr: hdr}:
		return nil
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
}

// send is a long running goroutine that sends data
func (s *Session) send() {
	for {
		select {
		case ready := <-s.sendCh:
			// Send a header if ready
			if ready.Hdr != nil {
				sent := 0
				for sent < len(ready.Hdr) {
					n, err := s.conn.Write(ready.Hdr[sent:])
					if err != nil {
						s.exitErr(err)
						asyncSendErr(ready.Err, err)
						return
					}
					sent += n
				}
			}

			// Send data from a body if given
			if ready.Body != nil {
				_, err := io.Copy(s.conn, ready.Body)
				if err != nil {
					s.exitErr(err)
					asyncSendErr(ready.Err, err)
					return
				}
			}

			// No error, successful send
			asyncSendErr(ready.Err, nil)
		case <-s.shutdownCh:
			return
		}
	}
}

// recv is a long running goroutine that accepts new data
func (s *Session) recv() {
	hdr := header(make([]byte, headerSize))
	for !s.isShutdown() {
		// Read the header
		if _, err := io.ReadFull(s.conn, hdr); err != nil {
			s.exitErr(err)
			return
		}

		// Verify the version
		if hdr.Version() != protoVersion {
			s.exitErr(ErrInvalidVersion)
			return
		}

		// Switch on the type
		msgType := hdr.MsgType()
		switch msgType {
		case typeData:
			fallthrough
		case typeWindowUpdate:
			if err := s.handleStreamMessage(hdr); err != nil {
				s.exitErr(err)
				return
			}
		case typeGoAway:
			if err := s.handleGoAway(hdr); err != nil {
				s.exitErr(err)
				return
			}
		case typePing:
			if err := s.handlePing(hdr); err != nil {
				s.exitErr(err)
				return
			}
		default:
			s.exitErr(ErrInvalidMsgType)
			return
		}
	}
}

// handleStreamMessage handles either a data or window update frame
func (s *Session) handleStreamMessage(hdr header) error {
	// Check for a new stream creation
	id := hdr.StreamID()
	flags := hdr.Flags()
	if flags&flagSYN == flagSYN {
		if err := s.incomingStream(id); err != nil {
			return err
		}
	}

	// Get the stream
	s.streamLock.RLock()
	stream := s.streams[id]
	s.streamLock.RUnlock()

	// Make sure we have a stream
	if stream == nil {
		s.goAway(goAwayProtoErr)
		return ErrMissingStream
	}

	// Check if this is a window update
	if hdr.MsgType() == typeWindowUpdate {
		if err := stream.incrSendWindow(hdr, flags); err != nil {
			s.goAway(goAwayProtoErr)
			return err
		}
	}

	// Read the new data
	if err := stream.readData(hdr, flags, s.conn); err != nil {
		s.goAway(goAwayProtoErr)
		return err
	}
	return nil
}

// handlePing is invokde for a typePing frame
func (s *Session) handlePing(hdr header) error {
	flags := hdr.Flags()
	pingID := hdr.Length()

	// Check if this is a query, respond back
	if flags&flagSYN == flagSYN {
		hdr := header(make([]byte, headerSize))
		hdr.encode(typePing, flagACK, 0, pingID)
		s.sendNoWait(hdr)
		return nil
	}

	// Handle a response
	s.pingLock.Lock()
	ch := s.pings[pingID]
	if ch != nil {
		delete(s.pings, pingID)
		close(ch)
	}
	s.pingLock.Unlock()
	return nil
}

// handleGoAway is invokde for a typeGoAway frame
func (s *Session) handleGoAway(hdr header) error {
	code := hdr.Length()
	switch code {
	case goAwayNormal:
		s.remoteGoAway = true
	case goAwayProtoErr:
		return fmt.Errorf("yamux protocol error")
	case goAwayInternalErr:
		return fmt.Errorf("remote yamux internal error")
	default:
		return fmt.Errorf("unexpected go away received")
	}
	return nil
}

// exitErr is used to handle an error that is causing
// the listener to exit.
func (s *Session) exitErr(err error) {
	s.shutdownErr = err
	s.Close()
}

// goAway is used to send a goAway message
func (s *Session) goAway(reason uint32) {
	hdr := header(make([]byte, headerSize))
	hdr.encode(typeGoAway, 0, 0, reason)
	s.sendNoWait(hdr)
}

// incomingStream is used to create a new incoming stream
func (s *Session) incomingStream(id uint32) error {
	// Reject immediately if we are doing a go away
	if s.localGoAway {
		hdr := header(make([]byte, headerSize))
		hdr.encode(typeWindowUpdate, flagRST, id, 0)
		s.sendNoWait(hdr)
		return nil
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()

	// Check if stream already exists
	if _, ok := s.streams[id]; ok {
		s.goAway(goAwayProtoErr)
		s.exitErr(ErrDuplicateStream)
		return nil
	}

	// Register the stream
	stream := newStream(s, id, streamSYNReceived)
	s.streams[id] = stream

	// Check if we've exceeded the backlog
	select {
	case s.acceptCh <- stream:
		return nil
	default:
		// Backlog exceeded! RST the stream
		delete(s.streams, id)
		stream.sendHdr.encode(typeWindowUpdate, flagRST, id, 0)
		s.sendNoWait(stream.sendHdr)
	}
	return nil
}

// closeStream is used to close a stream once both sides have
// issued a close.
func (s *Session) closeStream(id uint32, withLock bool) {
	if !withLock {
		s.streamLock.Lock()
		defer s.streamLock.Unlock()
	}
	delete(s.streams, id)
}
<commit_msg>Increase safety of GoAway<commit_after>package yamux

import (
	"fmt"
	"io"
	"math"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Session is used to wrap a reliable ordered connection and to
// multiplex it into multiple streams.
type Session struct {
	// remoteGoAway indicates the remote side does
	// not want futher connections. Must be first for alignment.
	remoteGoAway int32

	// localGoAway indicates that we should stop
	// accepting futher connections. Must be first for alignment.
	localGoAway int32

	// client is true if we are a client size connection
	client bool

	// config holds our configuration
	config *Config

	// conn is the underlying connection
	conn io.ReadWriteCloser

	// pings is used to track inflight pings
	pings    map[uint32]chan struct{}
	pingID   uint32
	pingLock sync.Mutex

	// nextStreamID is the next stream we should
	// send. This depends if we are a client/server.
	nextStreamID uint32

	// streams maps a stream id to a stream
	streams    map[uint32]*Stream
	streamLock sync.RWMutex

	// acceptCh is used to pass ready streams to the client
	acceptCh chan *Stream

	// sendCh is used to mark a stream as ready to send,
	// or to send a header out directly.
	sendCh chan sendReady

	// shutdown is used to safely close a session
	shutdown     bool
	shutdownErr  error
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// sendReady is used to either mark a stream as ready
// or to directly send a header
type sendReady struct {
	Hdr  []byte
	Body io.Reader
	Err  chan error
}

// newSession is used to construct a new session
func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := &Session{
		client:     client,
		config:     config,
		conn:       conn,
		pings:      make(map[uint32]chan struct{}),
		streams:    make(map[uint32]*Stream),
		acceptCh:   make(chan *Stream, config.AcceptBacklog),
		sendCh:     make(chan sendReady, 64),
		shutdownCh: make(chan struct{}),
	}
	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 2
	}
	go s.recv()
	go s.send()
	if config.EnableKeepAlive {
		go s.keepalive()
	}
	return s
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.shutdownCh:
		return true
	default:
		return false
	}
}

// Open is used to create a new stream
func (s *Session) Open() (*Stream, error) {
	if s.IsClosed() {
		return nil, ErrSessionShutdown
	}
	if atomic.LoadInt32(&s.remoteGoAway) == 1 {
		return nil, ErrRemoteGoAway
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()

	// Check if we've exhaused the streams
	id := s.nextStreamID
	if id >= math.MaxUint32-1 {
		return nil, ErrStreamsExhausted
	}
	s.nextStreamID += 2

	// Register the stream
	stream := newStream(s, id, streamInit)
	s.streams[id] = stream

	// Send the window update to create
	return stream, stream.sendWindowUpdate()
}

// Accept is used to block until the next available stream
// is ready to be accepted.
func (s *Session) Accept() (net.Conn, error) {
	return s.AcceptStream()
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	select {
	case stream := <-s.acceptCh:
		return stream, nil
	case <-s.shutdownCh:
		return nil, s.shutdownErr
	}
}

// Close is used to close the session and all streams.
// Attempts to send a GoAway before closing the connection.
func (s *Session) Close() error {
	s.shutdownLock.Lock()
	defer s.shutdownLock.Unlock()

	if s.shutdown {
		return nil
	}
	s.shutdown = true
	if s.shutdownErr == nil {
		s.shutdownErr = ErrSessionShutdown
	}
	close(s.shutdownCh)
	s.conn.Close()

	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	for _, stream := range s.streams {
		stream.forceClose()
	}
	return nil
}

// GoAway can be used to prevent accepting further
// connections. It does not close the underlying conn.
func (s *Session) GoAway() error {
	atomic.SwapInt32(&s.localGoAway, 1)
	s.goAway(goAwayNormal)
	return nil
}

// Ping is used to measure the RTT response time
func (s *Session) Ping() (time.Duration, error) {
	// Get a channel for the ping
	ch := make(chan struct{})

	// Get a new ping id, mark as pending
	s.pingLock.Lock()
	id := s.pingID
	s.pingID++
	s.pings[id] = ch
	s.pingLock.Unlock()

	// Send the ping request
	hdr := header(make([]byte, headerSize))
	hdr.encode(typePing, flagSYN, 0, id)
	if err := s.waitForSend(hdr, nil); err != nil {
		return 0, err
	}

	// Wait for a response
	start := time.Now()
	select {
	case <-ch:
	case <-s.shutdownCh:
		return 0, ErrSessionShutdown
	}

	// Compute the RTT
	return time.Now().Sub(start), nil
}

// keepalive is a long running goroutine that periodically does
// a ping to keep the connection alive.
func (s *Session) keepalive() {
	for {
		select {
		case <-time.After(s.config.KeepAliveInterval):
			s.Ping()
		case <-s.shutdownCh:
			return
		}
	}
}

// waitForSend waits to send a header, checking for a potential shutdown
func (s *Session) waitForSend(hdr header, body io.Reader) error {
	errCh := make(chan error, 1)
	ready := sendReady{Hdr: hdr, Body: body, Err: errCh}
	select {
	case s.sendCh <- ready:
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
	select {
	case err := <-errCh:
		return err
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
}

// sendNoWait does a send without waiting
func (s *Session) sendNoWait(hdr header) error {
	select {
	case s.sendCh <- sendReady{Hdr: hdr}:
		return nil
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
}

// send is a long running goroutine that sends data
func (s *Session) send() {
	for {
		select {
		case ready := <-s.sendCh:
			// Send a header if ready
			if ready.Hdr != nil {
				sent := 0
				for sent < len(ready.Hdr) {
					n, err := s.conn.Write(ready.Hdr[sent:])
					if err != nil {
						s.exitErr(err)
						asyncSendErr(ready.Err, err)
						return
					}
					sent += n
				}
			}

			// Send data from a body if given
			if ready.Body != nil {
				_, err := io.Copy(s.conn, ready.Body)
				if err != nil {
					s.exitErr(err)
					asyncSendErr(ready.Err, err)
					return
				}
			}

			// No error, successful send
			asyncSendErr(ready.Err, nil)
		case <-s.shutdownCh:
			return
		}
	}
}

// recv is a long running goroutine that accepts new data
func (s *Session) recv() {
	hdr := header(make([]byte, headerSize))
	for !s.IsClosed() {
		// Read the header
		if _, err := io.ReadFull(s.conn, hdr); err != nil {
			s.exitErr(err)
			return
		}

		// Verify the version
		if hdr.Version() != protoVersion {
			s.exitErr(ErrInvalidVersion)
			return
		}

		// Switch on the type
		msgType := hdr.MsgType()
		switch msgType {
		case typeData:
			fallthrough
		case typeWindowUpdate:
			if err := s.handleStreamMessage(hdr); err != nil {
				s.exitErr(err)
				return
			}
		case typeGoAway:
			if err := s.handleGoAway(hdr); err != nil {
				s.exitErr(err)
				return
			}
		case typePing:
			if err := s.handlePing(hdr); err != nil {
				s.exitErr(err)
				return
			}
		default:
			s.exitErr(ErrInvalidMsgType)
			return
		}
	}
}

// handleStreamMessage handles either a data or window update frame
func (s *Session) handleStreamMessage(hdr header) error {
	// Check for a new stream creation
	id := hdr.StreamID()
	flags := hdr.Flags()
	if flags&flagSYN == flagSYN {
		if err := s.incomingStream(id); err != nil {
			return err
		}
	}

	// Get the stream
	s.streamLock.RLock()
	stream := s.streams[id]
	s.streamLock.RUnlock()

	// Make sure we have a stream
	if stream == nil {
		s.goAway(goAwayProtoErr)
		return ErrMissingStream
	}

	// Check if this is a window update
	if hdr.MsgType() == typeWindowUpdate {
		if err := stream.incrSendWindow(hdr, flags); err != nil {
			s.goAway(goAwayProtoErr)
			return err
		}
	}

	// Read the new data
	if err := stream.readData(hdr, flags, s.conn); err != nil {
		s.goAway(goAwayProtoErr)
		return err
	}
	return nil
}

// handlePing is invokde for a typePing frame
func (s *Session) handlePing(hdr header) error {
	flags := hdr.Flags()
	pingID := hdr.Length()

	// Check if this is a query, respond back
	if flags&flagSYN == flagSYN {
		hdr := header(make([]byte, headerSize))
		hdr.encode(typePing, flagACK, 0, pingID)
		s.sendNoWait(hdr)
		return nil
	}

	// Handle a response
	s.pingLock.Lock()
	ch := s.pings[pingID]
	if ch != nil {
		delete(s.pings, pingID)
		close(ch)
	}
	s.pingLock.Unlock()
	return nil
}

// handleGoAway is invokde for a typeGoAway frame
func (s *Session) handleGoAway(hdr header) error {
	code := hdr.Length()
	switch code {
	case goAwayNormal:
		atomic.SwapInt32(&s.remoteGoAway, 1)
	case goAwayProtoErr:
		return fmt.Errorf("yamux protocol error")
	case goAwayInternalErr:
		return fmt.Errorf("remote yamux internal error")
	default:
		return fmt.Errorf("unexpected go away received")
	}
	return nil
}

// exitErr is used to handle an error that is causing
// the listener to exit.
func (s *Session) exitErr(err error) {
	s.shutdownErr = err
	s.Close()
}

// goAway is used to send a goAway message
func (s *Session) goAway(reason uint32) {
	hdr := header(make([]byte, headerSize))
	hdr.encode(typeGoAway, 0, 0, reason)
	s.sendNoWait(hdr)
}

// incomingStream is used to create a new incoming stream
func (s *Session) incomingStream(id uint32) error {
	// Reject immediately if we are doing a go away
	if atomic.LoadInt32(&s.localGoAway) == 1 {
		hdr := header(make([]byte, headerSize))
		hdr.encode(typeWindowUpdate, flagRST, id, 0)
		s.sendNoWait(hdr)
		return nil
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()

	// Check if stream already exists
	if _, ok := s.streams[id]; ok {
		s.goAway(goAwayProtoErr)
		s.exitErr(ErrDuplicateStream)
		return nil
	}

	// Register the stream
	stream := newStream(s, id, streamSYNReceived)
	s.streams[id] = stream

	// Check if we've exceeded the backlog
	select {
	case s.acceptCh <- stream:
		return nil
	default:
		// Backlog exceeded! RST the stream
		delete(s.streams, id)
		stream.sendHdr.encode(typeWindowUpdate, flagRST, id, 0)
		s.sendNoWait(stream.sendHdr)
	}
	return nil
}

// closeStream is used to close a stream once both sides have
// issued a close.
func (s *Session) closeStream(id uint32, withLock bool) {
	if !withLock {
		s.streamLock.Lock()
		defer s.streamLock.Unlock()
	}
	delete(s.streams, id)
}
<|endoftext|>


<commit_before>package revel

import (
	"fmt"
	"github.com/streadway/simpleuuid"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// A signed cookie (and thus limited to 4kb in size).
// Restriction: Keys may not have a colon in them.
type Session map[string]string

const (
	SESSION_ID_KEY = "_ID"
	TS_KEY         = "_TS"
)

// expireAfterDuration is the time to live, in seconds, of a session cookie.
// It may be specified in config as "session.expires". Values greater than 0
// set a persistent cookie with a time to live as specified, and the value 0
// sets a session cookie.
var expireAfterDuration time.Duration

func init() {
	// Set expireAfterDuration, default to 30 days if no value in config
	OnAppStart(func() {
		var err error
		if expiresString, ok := Config.String("session.expires"); !ok {
			expireAfterDuration = 30 * 24 * time.Hour
		} else if expiresString == "session" {
			expireAfterDuration = 0
		} else if expireAfterDuration, err = time.ParseDuration(expiresString); err != nil {
			panic(fmt.Errorf("session.expires invalid: %s", err))
		}
	})
}

// Id retrieves from the cookie or creates a time-based UUID identifying this
// session.
func (s Session) Id() string {
	if uuidStr, ok := s[SESSION_ID_KEY]; ok {
		return uuidStr
	}

	uuid, err := simpleuuid.NewTime(time.Now())
	if err != nil {
		panic(err) // I don't think this can actually happen.
	}
	s[SESSION_ID_KEY] = uuid.String()
	return s[SESSION_ID_KEY]
}

// getSessionExpiration return a time.Time with the session's expiration date
func getSessionExpiration() time.Time {
	if expireAfterDuration == 0 {
		return time.Time{}
	}
	return time.Now().Add(expireAfterDuration)
}

// cookie returns an http.Cookie containing the signed session.
func (s Session) cookie() *http.Cookie {
	var sessionValue string
	ts := getSessionExpiration()
	s[TS_KEY] = getSessionExpirationCookie(ts)
	for key, value := range s {
		if strings.ContainsAny(key, ":\x00") {
			panic("Session keys may not have colons or null bytes")
		}
		if strings.Contains(value, "\x00") {
			panic("Session values may not have null bytes")
		}
		sessionValue += "\x00" + key + ":" + value + "\x00"
	}

	sessionData := url.QueryEscape(sessionValue)
	return &http.Cookie{
		Name:     CookiePrefix + "_SESSION",
		Value:    Sign(sessionData) + "-" + sessionData,
		Path:     "/",
		HttpOnly: CookieHttpOnly,
		Secure:   CookieSecure,
		Expires:  ts.UTC(),
	}
}

// sessionTimeoutExpiredOrMissing returns a boolean of whether the session
// cookie is either not present or present but beyond its time to live; i.e.,
// whether there is not a valid session.
func sessionTimeoutExpiredOrMissing(session Session) bool {
	if exp, present := session[TS_KEY]; !present {
		return true
	} else if exp == "session" {
		return false
	} else if expInt, _ := strconv.Atoi(exp); int64(expInt) < time.Now().Unix() {
		return true
	}
	return false
}

// getSessionFromCookie returns a Session struct pulled from the signed
// session cookie.
func getSessionFromCookie(cookie *http.Cookie) Session {
	session := make(Session)

	// Separate the data from the signature.
	hyphen := strings.Index(cookie.Value, "-")
	if hyphen == -1 || hyphen >= len(cookie.Value)-1 {
		return session
	}
	sig, data := cookie.Value[:hyphen], cookie.Value[hyphen+1:]

	// Verify the signature.
	if !Verify(data, sig) {
		INFO.Println("Session cookie signature failed")
		return session
	}

	ParseKeyValueCookie(data, func(key, val string) {
		session[key] = val
	})

	if sessionTimeoutExpiredOrMissing(session) {
		session = make(Session)
	}

	return session
}

// SessionFilter is a Revel Filter that retrieves and sets the session cookie.
// Within Revel, it is available as a Session attribute on Controller instances.
// The name of the Session cookie is set as CookiePrefix + "_SESSION".
func SessionFilter(c *Controller, fc []Filter) {
	c.Session = restoreSession(c.Request.Request)
	// Make session vars available in templates as {{.session.xyz}}
	c.RenderArgs["session"] = c.Session

	fc[0](c, fc[1:])

	// Store the session (and sign it).
	c.SetCookie(c.Session.cookie())
}

// restoreSession returns either the current session, retrieved from the
// session cookie, or a new session.
func restoreSession(req *http.Request) Session {
	cookie, err := req.Cookie(CookiePrefix + "_SESSION")
	if err != nil {
		return make(Session)
	} else {
		return getSessionFromCookie(cookie)
	}
}

// getSessionExpirationCookie retrieves the cookie's time to live as a
// string of either the number of seconds, for a persistent cookie, or
// "session".
func getSessionExpirationCookie(t time.Time) string {
	if t.IsZero() {
		return "session"
	}
	return strconv.FormatInt(t.Unix(), 10)
}
<commit_msg>session expire for per session<commit_after>package revel

import (
	"fmt"
	"github.com/streadway/simpleuuid"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// A signed cookie (and thus limited to 4kb in size).
// Restriction: Keys may not have a colon in them.
type Session map[string]string

const (
	SESSION_ID_KEY = "_ID"
	TS_KEY         = "_TS"
)

// expireAfterDuration is the time to live, in seconds, of a session cookie.
// It may be specified in config as "session.expires". Values greater than 0
// set a persistent cookie with a time to live as specified, and the value 0
// sets a session cookie.
var expireAfterDuration time.Duration

func init() {
	// Set expireAfterDuration, default to 30 days if no value in config
	OnAppStart(func() {
		var err error
		if expiresString, ok := Config.String("session.expires"); !ok {
			expireAfterDuration = 30 * 24 * time.Hour
		} else if expiresString == "session" {
			expireAfterDuration = 0
		} else if expireAfterDuration, err = time.ParseDuration(expiresString); err != nil {
			panic(fmt.Errorf("session.expires invalid: %s", err))
		}
	})
}

// Id retrieves from the cookie or creates a time-based UUID identifying this
// session.
func (s Session) Id() string {
	if uuidStr, ok := s[SESSION_ID_KEY]; ok {
		return uuidStr
	}

	uuid, err := simpleuuid.NewTime(time.Now())
	if err != nil {
		panic(err) // I don't think this can actually happen.
	}
	s[SESSION_ID_KEY] = uuid.String()
	return s[SESSION_ID_KEY]
}

// getSessionExpiration return a time.Time with the session's expiration date.
// If previous session has set to "session", remain it
func (s Session) getSessionExpiration() time.Time {
	if expireAfterDuration == 0 || s[TS_KEY] == "session" {
		return time.Time{}
	}
	return time.Now().Add(expireAfterDuration)
}

// cookie returns an http.Cookie containing the signed session.
func (s Session) cookie() *http.Cookie {
	var sessionValue string
	ts := s.getSessionExpiration()
	s[TS_KEY] = getSessionExpirationCookie(ts)
	for key, value := range s {
		if strings.ContainsAny(key, ":\x00") {
			panic("Session keys may not have colons or null bytes")
		}
		if strings.Contains(value, "\x00") {
			panic("Session values may not have null bytes")
		}
		sessionValue += "\x00" + key + ":" + value + "\x00"
	}

	sessionData := url.QueryEscape(sessionValue)
	return &http.Cookie{
		Name:     CookiePrefix + "_SESSION",
		Value:    Sign(sessionData) + "-" + sessionData,
		Path:     "/",
		HttpOnly: CookieHttpOnly,
		Secure:   CookieSecure,
		Expires:  ts.UTC(),
	}
}

// sessionTimeoutExpiredOrMissing returns a boolean of whether the session
// cookie is either not present or present but beyond its time to live; i.e.,
// whether there is not a valid session.
func sessionTimeoutExpiredOrMissing(session Session) bool {
	if exp, present := session[TS_KEY]; !present {
		return true
	} else if exp == "session" {
		return false
	} else if expInt, _ := strconv.Atoi(exp); int64(expInt) < time.Now().Unix() {
		return true
	}
	return false
}

// getSessionFromCookie returns a Session struct pulled from the signed
// session cookie.
func getSessionFromCookie(cookie *http.Cookie) Session {
	session := make(Session)

	// Separate the data from the signature.
	hyphen := strings.Index(cookie.Value, "-")
	if hyphen == -1 || hyphen >= len(cookie.Value)-1 {
		return session
	}
	sig, data := cookie.Value[:hyphen], cookie.Value[hyphen+1:]

	// Verify the signature.
	if !Verify(data, sig) {
		INFO.Println("Session cookie signature failed")
		return session
	}

	ParseKeyValueCookie(data, func(key, val string) {
		session[key] = val
	})

	if sessionTimeoutExpiredOrMissing(session) {
		session = make(Session)
	}

	return session
}

// SessionFilter is a Revel Filter that retrieves and sets the session cookie.
// Within Revel, it is available as a Session attribute on Controller instances.
// The name of the Session cookie is set as CookiePrefix + "_SESSION".
func SessionFilter(c *Controller, fc []Filter) {
	c.Session = restoreSession(c.Request.Request)
	// Make session vars available in templates as {{.session.xyz}}
	c.RenderArgs["session"] = c.Session

	fc[0](c, fc[1:])

	// Store the session (and sign it).
	c.SetCookie(c.Session.cookie())
}

// restoreSession returns either the current session, retrieved from the
// session cookie, or a new session.
func restoreSession(req *http.Request) Session {
	cookie, err := req.Cookie(CookiePrefix + "_SESSION")
	if err != nil {
		return make(Session)
	} else {
		return getSessionFromCookie(cookie)
	}
}

// getSessionExpirationCookie retrieves the cookie's time to live as a
// string of either the number of seconds, for a persistent cookie, or
// "session".
func getSessionExpirationCookie(t time.Time) string {
	if t.IsZero() {
		return "session"
	}
	return strconv.FormatInt(t.Unix(), 10)
}

// set session expiration to "session"
func (s Session) SetSessionNoExpiration() {
	s[TS_KEY] = "session"
}

// reset the session expiration to default
func (s Session) SetSessionDefaultExpiration() {
	delete(TS_KEY)
}
<|endoftext|>