server.go

package web

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/pprof"
	"os"
	"path"
	"reflect"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"

	"code.google.com/p/go.net/websocket"
)

// ServerConfig is configuration for server objects.
type ServerConfig struct {
	StaticHost            string
	StaticHostExcludeType string
	StaticHostExcludeFile string
	StaticDir             string
	Addr                  string
	Port                  int
	CookieSecret          string
	XSRFSecret            string
	RecoverPanic          bool
	Profiler              bool
}

// Server represents a web.go server.
type Server struct {
	Config         *ServerConfig
	routes         []route
	Logger         *log.Logger
	Env            map[string]interface{}
	SessionStorage ISessionStorage
	XSRFGetUid     func(*Context) string
	//save the listener so it can be closed
	l                net.Listener
	enableXSRF       bool
	ErrorHandler     func(errorMsg string)
	ErrorPageHandler func(*Context, int, interface{}) string
}

func NewServer() *Server {
	return &Server{
		Config:         &ServerConfig{},
		Logger:         log.New(os.Stdout, "", log.Ldate|log.Ltime),
		Env:            map[string]interface{}{},
		SessionStorage: MemoryStore,
	}
}

func (s *Server) initServer() {
	if s.Config == nil {
		s.Config = &ServerConfig{}
	}

	if s.Logger == nil {
		s.Logger = log.New(os.Stdout, "", log.Ldate|log.Ltime)
	}
}

type route struct {
	r           string
	cr          *regexp.Regexp
	method      string
	handler     reflect.Value
	httpHandler http.Handler
}

func (s *Server) addRoute(r string, method string, handler interface{}) {
	cr, err := regexp.Compile(r)
	if err != nil {
		s.Logger.Printf("Error in route regex %q\n", r)
		return
	}

	switch handler.(type) {
	case http.Handler:
		s.routes = append(s.routes, route{r: r, cr: cr, method: method, httpHandler: handler.(http.Handler)})
	case reflect.Value:
		fv := handler.(reflect.Value)
		s.routes = append(s.routes, route{r: r, cr: cr, method: method, handler: fv})
	default:
		fv := reflect.ValueOf(handler)
		s.routes = append(s.routes, route{r: r, cr: cr, method: method, handler: fv})
	}
}

// ServeHTTP is the interface method for Go's http server package
func (s *Server) ServeHTTP(c http.ResponseWriter, req *http.Request) {
	s.Process(c, req)
}

// Process invokes the routing system for server s
func (s *Server) Process(c http.ResponseWriter, req *http.Request) {
	route := s.routeHandler(req, c)
	if route != nil {
		route.httpHandler.ServeHTTP(c, req)
	}
}

// Get adds a handler for the 'GET' http method for server s.
func (s *Server) Get(route string, handler interface{}) {
	s.addRoute(route, "GET", handler)
}

// Post adds a handler for the 'POST' http method for server s.
func (s *Server) Post(route string, handler interface{}) {
	s.addRoute(route, "POST", handler)
}

// Put adds a handler for the 'PUT' http method for server s.
func (s *Server) Put(route string, handler interface{}) {
	s.addRoute(route, "PUT", handler)
}

// Delete adds a handler for the 'DELETE' http method for server s.
func (s *Server) Delete(route string, handler interface{}) {
	s.addRoute(route, "DELETE", handler)
}

// Match adds a handler for an arbitrary http method for server s.
func (s *Server) Match(method string, route string, handler interface{}) {
	s.addRoute(route, method, handler)
}

//Adds a custom handler. Only for webserver mode. Will have no effect when running as FCGI or SCGI.
func (s *Server) Handler(route string, method string, httpHandler http.Handler) {
	s.addRoute(route, method, httpHandler)
}

//Adds a handler for websockets. Only for webserver mode. Will have no effect when running as FCGI or SCGI.
func (s *Server) Websocket(route string, httpHandler websocket.Handler) {
	s.addRoute(route, "GET", httpHandler)
}

// Run starts the web application and serves HTTP requests for s
func (s *Server) Run(addr string) {
	s.initServer()

	mux := http.NewServeMux()
	if s.Config.Profiler {
		mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
		mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
		mux.Handle("/debug/pprof/heap", pprof.Handler("heap"))
		mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
	}
	mux.Handle("/", s)

	s.Logger.Printf("web.go serving %s\n", addr)

	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("ListenAndServe:", err)
	}
	s.l = l
	err = http.Serve(s.l, mux)
	s.l.Close()
}

// RunFcgi starts the web application and serves FastCGI requests for s.
func (s *Server) RunFcgi(addr string) {
	s.initServer()
	s.Logger.Printf("web.go serving fcgi %s\n", addr)
	s.listenAndServeFcgi(addr)
}

// RunScgi starts the web application and serves SCGI requests for s.
func (s *Server) RunScgi(addr string) {
	s.initServer()
	s.Logger.Printf("web.go serving scgi %s\n", addr)
	s.listenAndServeScgi(addr)
}

// RunTLS starts the web application and serves HTTPS requests for s.
func (s *Server) RunTLS(addr string, config *tls.Config) error {
	s.initServer()
	mux := http.NewServeMux()
	mux.Handle("/", s)
	l, err := tls.Listen("tcp", addr, config)
	if err != nil {
		log.Fatal("Listen:", err)
		return err
	}

	s.l = l
	return http.Serve(s.l, mux)
}

// Close stops server s.
func (s *Server) Close() {
	if s.l != nil {
		s.l.Close()
	}
}

// safelyCall invokes `function` in recover block
func (s *Server) safelyCall(req *http.Request, function reflect.Value, args []reflect.Value) (resp []reflect.Value, e interface{}, stackDump string) {
	defer func() {
		if err := recover(); err != nil {
			if !s.Config.RecoverPanic {
				// go back to panic
				panic(err)
			} else {
				e = err
				resp = nil
				var errors []string

				for i := 1; ; i += 1 {
					_, file, line, ok := runtime.Caller(i)
					if !ok {
						break
					}
					errors = append(errors, fmt.Sprintf("\t%s:%d", file, line))
				}

				errorMsg := req.Method + " " + req.RequestURI + "\n\n" + strings.Join(errors, "\n")
				stackDump = errorMsg

				s.Logger.Println("Handler crashed with error: ", err, "\n", errorMsg)
				if s.ErrorHandler != nil {
					errorMsg = fmt.Sprintf("Handler crashed with error: %v\n", err) + errorMsg
					s.ErrorHandler(errorMsg)
				}
			}
		}
	}()
	return function.Call(args), nil, ""
}

// requiresContext determines whether 'handlerType' contains
// an argument to 'web.Ctx' as its first argument
func requiresContext(handlerType reflect.Type) bool {
	//if the method doesn't take arguments, no
	if handlerType.NumIn() == 0 {
		return false
	}

	//if the first argument is not a pointer, no
	a0 := handlerType.In(0)
	if a0.Kind() != reflect.Ptr {
		return false
	}
	//if the first argument is a context, yes
	if a0.Elem() == contextType {
		return true
	}

	return false
}

func requiresIContext(handlerType reflect.Type) (IContext, bool) {
	//if the method doesn't take arguments, no
	if handlerType.NumIn() == 0 {
		return nil, false
	}

	//if the first argument is not a pointer, no
	a0 := handlerType.In(0)
	if a0.Kind() != reflect.Ptr {
		return nil, false
	}
	//if the first argument implements IContext, yes
	arg := reflect.New(a0).Elem().Interface()
	if ctx, ok := arg.(IContext); ok {
		return ctx, true
	} else {
		print("IContext test failed\n")
	}

	return nil, false
}

// tryServingFile attempts to serve a static file, and returns
// whether or not the operation is successful.
// It checks the following directories for the file, in order:
// 1) Config.StaticDir
// 2) The 'static' directory in the parent directory of the executable.
// 3) The 'static' directory in the current working directory
func (s *Server) tryServingFile(name string, req *http.Request, w http.ResponseWriter) bool {
	//try to serve a static file
	if s.Config.StaticDir != "" {
		staticFile := path.Join(s.Config.StaticDir, name)
		if fileExists(staticFile) {
			http.ServeFile(w, req, staticFile)
			return true
		}
	} else {
		for _, staticDir := range defaultStaticDirs {
			staticFile := path.Join(staticDir, name)
			if fileExists(staticFile) {
				http.ServeFile(w, req, staticFile)
				return true
			}
		}
	}
	return false
}

func (s *Server) logRequest(ctx Context, sTime time.Time) {
	//log the request
	var logEntry bytes.Buffer
	req := ctx.Request
	requestPath := req.URL.Path

	duration := time.Now().Sub(sTime)
	var client string

	// We suppose RemoteAddr is of the form Ip:Port as specified in the Request
	// documentation at http://golang.org/pkg/net/http/#Request
	pos := strings.LastIndex(req.RemoteAddr, ":")
	if pos > 0 {
		client = req.RemoteAddr[0:pos]
	} else {
		client = req.RemoteAddr
	}

	fmt.Fprintf(&logEntry, "%s - \033[32;1m %s %s\033[0m - %v", client, req.Method, requestPath, duration)

	if len(ctx.Params) > 0 {
		fmt.Fprintf(&logEntry, " - \033[37;1mParams: %v\033[0m\n", ctx.Params)
	}

	ctx.Server.Logger.Print(logEntry.String())

}

// the main route handler in web.go
// Tries to handle the given request.
// Finds the route matching the request, and execute the callback associated
// with it.  In case of custom http handlers, this function returns an "unused"
// route. The caller is then responsible for calling the httpHandler associated
// with the returned route.
func (s *Server) routeHandler(req *http.Request, w http.ResponseWriter) (unused *route) {
	requestPath := req.URL.Path
	ctx := Context{
		Request:        req,
		Params:         map[string]string{},
		Server:         s,
		ResponseWriter: w,
	}

	//set some default headers
	ctx.SetHeader("Server", "web.go", true)
	tm := time.Now().UTC()

	//ignore errors from ParseForm because it's usually harmless.
	req.ParseForm()
	if len(req.Form) > 0 {
		for k, v := range req.Form {
			ctx.Params[k] = v[0]
		}
	}

	defer s.logRequest(ctx, tm)

	ctx.SetHeader("Date", webTime(tm), true)

	if req.Method == "GET" || req.Method == "HEAD" {
		if s.tryServingFile(requestPath, req, w) {
			return
		}
	}

	//Set the default content-type
	ctx.SetHeader("Content-Type", "text/html; charset=utf-8", true)

	for i := 0; i < len(s.routes); i++ {
		route := s.routes[i]
		cr := route.cr
		//if the methods don't match, skip this handler (except HEAD can be used in place of GET)
		if req.Method != route.method && !(req.Method == "HEAD" && route.method == "GET") {
			continue
		}

		if !cr.MatchString(requestPath) {
			continue
		}
		match := cr.FindStringSubmatch(requestPath)

		if len(match[0]) != len(requestPath) {
			continue
		}

		if route.httpHandler != nil {
			unused = &route
			// We can not handle custom http handlers here, give back to the caller.
			return
		}

		// generate or get XSRF token
		if s.enableXSRF {
			getXSRFToken(s, &ctx)
			if req.Method == "POST" && !XSRFValidate(&ctx) {
				ctx.ResponseWriter.Write([]byte("Invalid XSRF"))
				ctx.Forbidden()
				return
			}
		}

		var args []reflect.Value
		handlerType := route.handler.Type()
		if requiresContext(handlerType) {
			args = append(args, reflect.ValueOf(&ctx))
		}

		if newContext, ok := requiresIContext(handlerType); ok {
			obj := newContext.Create(&ctx)

			//todo: must catch panic
			continueToHandler := obj.BeforeHandler()
			if continueToHandler != true {
				return
			}
			args = append(args, reflect.ValueOf(obj))
		}

		for _, arg := range match[1:] {
			args = append(args, reflect.ValueOf(arg))
		}

		ret, err, dump := s.safelyCall(req, route.handler, args)
		if err != nil {
			//there was an error or panic while calling the handler
			es := "Server Error"
			if s.ErrorPageHandler != nil {
				args := map[string]interface{}{"error": err, "dump": dump}
				es = s.ErrorPageHandler(&ctx, 500, args)
			}
			ctx.Abort(500, es)
		}
		if len(ret) == 0 {
			return
		}

		sval := ret[0]

		var content []byte

		if sval.Kind() == reflect.String {
			content = []byte(sval.String())
		} else if sval.Kind() == reflect.Slice && sval.Type().Elem().Kind() == reflect.Uint8 {
			content = sval.Interface().([]byte)
		}
		ctx.SetHeader("Content-Length", strconv.Itoa(len(content)), true)
		_, err = ctx.ResponseWriter.Write(content)
		if err != nil {
			ctx.Server.Logger.Println("Error during write: ", err)
		}
		return
	}

	// try serving index.html or index.htm
	if req.Method == "GET" || req.Method == "HEAD" {
		if s.tryServingFile(path.Join(requestPath, "index.html"), req, w) {
			return
		} else if s.tryServingFile(path.Join(requestPath, "index.htm"), req, w) {
			return
		}
	}
	es := "Page not found"
	if s.ErrorPageHandler != nil {
		es = s.ErrorPageHandler(&ctx, 404, nil)
	}
	ctx.Abort(404, es)
	return
}

// SetLogger sets the logger for server s
func (s *Server) SetLogger(logger *log.Logger) {
	s.Logger = logger
}

func (s *Server) SetXSRFOption(getUid func(*Context) string) {
	s.XSRFGetUid = getUid
	s.enableXSRF = true
}

func (s *Server) SetSessionStorage(ss ISessionStorage) {
	s.SessionStorage = ss
}

func (s *Server) SetErrorHandler(f func(errorMsg string)) {
	s.ErrorHandler = f
}

func (s *Server) SetErrorPageHandler(f func(*Context, int, interface{}) string) {
	s.ErrorPageHandler = f
}