<commit_before>package dotweb

import (
	"fmt"
	"github.com/devfeel/dotweb/framework/convert"
	"github.com/devfeel/dotweb/framework/exception"
	"github.com/devfeel/dotweb/framework/json"
	"github.com/devfeel/dotweb/framework/log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/julienschmidt/httprouter"
	"golang.org/x/net/context/ctxhttp"
	"golang.org/x/net/websocket"
)

const (
	RouteMethod_GET       = "GET"
	RouteMethod_HEAD      = "HEAD"
	RouteMethod_OPTIONS   = "OPTIONS"
	RouteMethod_POST      = "POST"
	RouteMethod_PUT       = "PUT"
	RouteMethod_PATCH     = "PATCH"
	RouteMethod_DELETE    = "DELETE"
	RouteMethod_HiJack    = "HiJack"
	RouteMethod_WebSocket = "WebSocket"
)

type (
	//HttpModule定义
	HttpModule struct {
		//响应请求时作为 HTTP 执行管线链中的第一个事件发生
		OnBeginRequest func(*HttpContext)
		//响应请求时作为 HTTP 执行管线链中的最后一个事件发生。
		OnEndRequest func(*HttpContext)
	}

	//HttpServer定义
	HttpServer struct {
		router *httprouter.Router
		dotweb *Dotweb
		pool   *pool
	}

	//pool定义
	pool struct {
		response sync.Pool
		context  sync.Pool
	}
)

// Handle is a function that can be registered to a route to handle HTTP
// requests. Like http.HandlerFunc, but has a third parameter for the values of
// wildcards (variables).
type HttpHandle func(*HttpContext)

func NewHttpServer() *HttpServer {
	server := &HttpServer{
		router: httprouter.New(),
		pool: &pool{
			response: sync.Pool{
				New: func() interface{} {
					return &Response{}
				},
			},
			context: sync.Pool{
				New: func() interface{} {
					return &HttpContext{}
				},
			},
		},
	}
	return server
}

//ServeHTTP makes the httprouter implement the http.Handler interface.
func (server *HttpServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	//针对websocket与调试信息特殊处理
	if checkIsWebSocketRequest(req) {
		http.DefaultServeMux.ServeHTTP(w, req)
	} else {
		//设置header信息
		w.Header().Set(HeaderServer, DefaultServerName)
		server.router.ServeHTTP(w, req)
	}
}

/*
* 关联当前HttpServer实例对应的DotServer实例
 */
func (server *HttpServer) setDotweb(dotweb *Dotweb) {
	server.dotweb = dotweb
}

// GET is a shortcut for router.Handle("GET", path, handle)
func (server *HttpServer) GET(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_GET, path, handle)
}

// HEAD is a shortcut for router.Handle("HEAD", path, handle)
func (server *HttpServer) HEAD(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_HEAD, path, handle)
}

// OPTIONS is a shortcut for router.Handle("OPTIONS", path, handle)
func (server *HttpServer) OPTIONS(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_OPTIONS, path, handle)
}

// POST is a shortcut for router.Handle("POST", path, handle)
func (server *HttpServer) POST(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_POST, path, handle)
}

// PUT is a shortcut for router.Handle("PUT", path, handle)
func (server *HttpServer) PUT(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_PUT, path, handle)
}

// PATCH is a shortcut for router.Handle("PATCH", path, handle)
func (server *HttpServer) PATCH(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_PATCH, path, handle)
}

// DELETE is a shortcut for router.Handle("DELETE", path, handle)
func (server *HttpServer) DELETE(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_DELETE, path, handle)
}

// DELETE is a shortcut for router.Handle("DELETE", path, handle)
func (server *HttpServer) HiJack(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_GET, path, handle)
}

// shortcut for router.Handle(httpmethod, path, handle)
// support GET\POST\DELETE\PUT\HEAD\PATCH\OPTIONS\HiJack\WebSocket
func (server *HttpServer) RegisterRoute(routeMethod string, path string, handle HttpHandle) {

	logger.Log("Dotweb:RegisterRoute ["+routeMethod+"] ["+path+"] ["+fmt.Sprintln(handle)+"]", LogTarget_HttpServer, LogLevel_Debug)

	//hijack mode,use get and isHijack = true
	if routeMethod == RouteMethod_HiJack {
		server.router.Handle(RouteMethod_GET, path, server.wrapRouterHandle(handle, true))
		return
	}
	//websocket mode,use default httpserver
	if routeMethod == RouteMethod_WebSocket {
		http.Handle(path, websocket.Handler(server.wrapWebSocketHandle(handle)))
		return
	}

	//GET\POST\DELETE\PUT\HEAD\PATCH\OPTIONS mode
	server.router.Handle(routeMethod, path, server.wrapRouterHandle(handle, false))
	return
}

// ServerFile is a shortcut for router.ServeFiles(path, filepath)
// simple demo:server.ServerFile("/src/*filepath", "/var/www")
func (server *HttpServer) ServerFile(urlpath string, filepath string) {
	server.router.ServeFiles(urlpath, http.Dir(filepath))
}

// WebSocket is a shortcut for websocket.Handler
func (server *HttpServer) WebSocket(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_WebSocket, path, handle)
}

type LogJson struct {
	RequestUrl string
	HttpHeader string
	HttpBody   string
}

//wrap HttpHandle to httprouter.Handle
func (server *HttpServer) wrapRouterHandle(handle HttpHandle, isHijack bool) httprouter.Handle {
	return func(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
		//get from pool
		res := server.pool.response.Get().(*Response)
		res.Reset(w)
		httpCtx := server.pool.context.Get().(*HttpContext)
		httpCtx.Reset(res, r, params)

		//增加状态计数
		GlobalState.AddRequestCount(1)

		//hijack处理
		if isHijack {
			_, hijack_err := httpCtx.Hijack()
			if hijack_err != nil {
				//输出内容
				httpCtx.Response.WriteHeader(http.StatusInternalServerError)
				httpCtx.Response.Header().Set(HeaderContentType, CharsetUTF8)
				httpCtx.WriteString(hijack_err.Error())
			}
		}

		startTime := time.Now()
		defer func() {
			var errmsg string
			if err := recover(); err != nil {
				errmsg = exception.CatchError("httpserver::RouterHandle", LogTarget_HttpServer, err)

				//默认异常处理
				if server.dotweb.ExceptionHandler != nil {
					server.dotweb.ExceptionHandler(httpCtx, err)
				}

				//记录访问日志
				headinfo := fmt.Sprintln(httpCtx.Response.Header)
				logJson := LogJson{
					RequestUrl: httpCtx.Request.RequestURI,
					HttpHeader: headinfo,
					HttpBody:   errmsg,
				}
				logString := jsonutil.GetJsonString(logJson)
				logger.Log(logString, LogTarget_HttpServer, LogLevel_Error)

				//增加错误计数
				GlobalState.AddErrorCount(1)
			}
			timetaken := int64(time.Now().Sub(startTime) / time.Millisecond)
			//HttpServer Logging
			logger.Log(httpCtx.Url()+" "+logString(httpCtx, timetaken), LogTarget_HttpRequest, LogLevel_Debug)

			// Return to pool
			server.pool.response.Put(res)
			server.pool.context.Put(httpCtx)
		}()

		//处理前置Module集合
		for _, module := range server.dotweb.Modules {
			if module.OnBeginRequest != nil {
				module.OnBeginRequest(httpCtx)
				//if current module set HttpContext.End,stop this request
				if httpCtx.isEnd {
					return
				}
			}
		}

		//处理用户handle
		handle(httpCtx)

		//处理后置Module集合
		for _, module := range server.dotweb.Modules {
			if module.OnEndRequest != nil {
				module.OnEndRequest(httpCtx)
				if httpCtx.isEnd {
					//if current module set HttpContext.End,stop this request
					return
				}
			}
		}

	}
}

//wrap HttpHandle to websocket.Handle
func (server *HttpServer) wrapWebSocketHandle(handle HttpHandle) websocket.Handler {
	return func(ws *websocket.Conn) {
		//get from pool
		httpCtx := server.pool.context.Get().(*HttpContext)
		httpCtx.Reset(nil, ws.Request(), nil)
		httpCtx.WebSocket = &WebSocket{
			Conn: ws,
		}
		httpCtx.IsWebSocket = true

		startTime := time.Now()
		defer func() {
			var errmsg string
			if err := recover(); err != nil {
				errmsg = exception.CatchError("httpserver::WebsocketHandle", LogTarget_HttpServer, err)

				//记录访问日志
				headinfo := fmt.Sprintln(httpCtx.WebSocket.Conn.Request().Header)
				logJson := LogJson{
					RequestUrl: httpCtx.WebSocket.Conn.Request().RequestURI,
					HttpHeader: headinfo,
					HttpBody:   errmsg,
				}
				logString := jsonutil.GetJsonString(logJson)
				logger.Log(logString, LogTarget_HttpServer, LogLevel_Error)

				//增加错误计数
				GlobalState.AddErrorCount(1)
			}
			timetaken := int64(time.Now().Sub(startTime) / time.Millisecond)
			//HttpServer Logging
			logger.Log(httpCtx.Url()+" "+logString(httpCtx, timetaken), LogTarget_HttpRequest, LogLevel_Debug)

			// Return to pool
			server.pool.context.Put(httpCtx)
		}()

		handle(httpCtx)

		//增加状态计数
		GlobalState.AddRequestCount(1)
	}
}

//get default log string
func logString(ctx *HttpContext, timetaken int64) string {

	var reqbytelen, resbytelen, method, proto, status, userip string
	if !ctx.IsWebSocket {
		reqbytelen = convert.Int642String(ctx.Request.ContentLength)
		resbytelen = convert.Int642String(ctx.Response.Size)
		method = ctx.Request.Method
		proto = ctx.Request.Proto
		status = convert.Int2String(ctx.Response.Status)
		userip = ctx.RemoteIP()
	} else {
		reqbytelen = convert.Int642String(ctx.Request.ContentLength)
		resbytelen = "0"
		method = ctx.Request.Method
		proto = ctx.Request.Proto
		status = "0"
		userip = ctx.RemoteIP()
	}

	log := method + " "
	log += userip + " "
	log += proto + " "
	log += status + " "
	log += reqbytelen + " "
	log += resbytelen + " "
	log += convert.Int642String(timetaken)

	return log
}

//check request is the websocket request
func checkIsWebSocketRequest(req *http.Request) bool {
	if req.Header.Get("Connection") == "Upgrade" {
		return true
	}
	return false
}

//check request is startwith /debug/
func checkIsDebugRequest(req *http.Request) bool {
	if strings.Index(req.RequestURI, "/debug/") == 0 {
		return true
	}
	return false
}
<commit_msg>add HttpContext.End(), it can stop current request<commit_after>package dotweb

import (
	"fmt"
	"github.com/devfeel/dotweb/framework/convert"
	"github.com/devfeel/dotweb/framework/exception"
	"github.com/devfeel/dotweb/framework/json"
	"github.com/devfeel/dotweb/framework/log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/julienschmidt/httprouter"
	"golang.org/x/net/context/ctxhttp"
	"golang.org/x/net/websocket"
)

const (
	RouteMethod_GET       = "GET"
	RouteMethod_HEAD      = "HEAD"
	RouteMethod_OPTIONS   = "OPTIONS"
	RouteMethod_POST      = "POST"
	RouteMethod_PUT       = "PUT"
	RouteMethod_PATCH     = "PATCH"
	RouteMethod_DELETE    = "DELETE"
	RouteMethod_HiJack    = "HiJack"
	RouteMethod_WebSocket = "WebSocket"
)

type (
	//HttpModule定义
	HttpModule struct {
		//响应请求时作为 HTTP 执行管线链中的第一个事件发生
		OnBeginRequest func(*HttpContext)
		//响应请求时作为 HTTP 执行管线链中的最后一个事件发生。
		OnEndRequest func(*HttpContext)
	}

	//HttpServer定义
	HttpServer struct {
		router *httprouter.Router
		dotweb *Dotweb
		pool   *pool
	}

	//pool定义
	pool struct {
		response sync.Pool
		context  sync.Pool
	}
)

// Handle is a function that can be registered to a route to handle HTTP
// requests. Like http.HandlerFunc, but has a third parameter for the values of
// wildcards (variables).
type HttpHandle func(*HttpContext)

func NewHttpServer() *HttpServer {
	server := &HttpServer{
		router: httprouter.New(),
		pool: &pool{
			response: sync.Pool{
				New: func() interface{} {
					return &Response{}
				},
			},
			context: sync.Pool{
				New: func() interface{} {
					return &HttpContext{}
				},
			},
		},
	}
	return server
}

//ServeHTTP makes the httprouter implement the http.Handler interface.
func (server *HttpServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	//针对websocket与调试信息特殊处理
	if checkIsWebSocketRequest(req) {
		http.DefaultServeMux.ServeHTTP(w, req)
	} else {
		//设置header信息
		w.Header().Set(HeaderServer, DefaultServerName)
		server.router.ServeHTTP(w, req)
	}
}

/*
* 关联当前HttpServer实例对应的DotServer实例
 */
func (server *HttpServer) setDotweb(dotweb *Dotweb) {
	server.dotweb = dotweb
}

// GET is a shortcut for router.Handle("GET", path, handle)
func (server *HttpServer) GET(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_GET, path, handle)
}

// HEAD is a shortcut for router.Handle("HEAD", path, handle)
func (server *HttpServer) HEAD(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_HEAD, path, handle)
}

// OPTIONS is a shortcut for router.Handle("OPTIONS", path, handle)
func (server *HttpServer) OPTIONS(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_OPTIONS, path, handle)
}

// POST is a shortcut for router.Handle("POST", path, handle)
func (server *HttpServer) POST(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_POST, path, handle)
}

// PUT is a shortcut for router.Handle("PUT", path, handle)
func (server *HttpServer) PUT(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_PUT, path, handle)
}

// PATCH is a shortcut for router.Handle("PATCH", path, handle)
func (server *HttpServer) PATCH(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_PATCH, path, handle)
}

// DELETE is a shortcut for router.Handle("DELETE", path, handle)
func (server *HttpServer) DELETE(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_DELETE, path, handle)
}

// DELETE is a shortcut for router.Handle("DELETE", path, handle)
func (server *HttpServer) HiJack(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_GET, path, handle)
}

// shortcut for router.Handle(httpmethod, path, handle)
// support GET\POST\DELETE\PUT\HEAD\PATCH\OPTIONS\HiJack\WebSocket
func (server *HttpServer) RegisterRoute(routeMethod string, path string, handle HttpHandle) {

	logger.Log("Dotweb:RegisterRoute ["+routeMethod+"] ["+path+"] ["+fmt.Sprintln(handle)+"]", LogTarget_HttpServer, LogLevel_Debug)

	//hijack mode,use get and isHijack = true
	if routeMethod == RouteMethod_HiJack {
		server.router.Handle(RouteMethod_GET, path, server.wrapRouterHandle(handle, true))
		return
	}
	//websocket mode,use default httpserver
	if routeMethod == RouteMethod_WebSocket {
		http.Handle(path, websocket.Handler(server.wrapWebSocketHandle(handle)))
		return
	}

	//GET\POST\DELETE\PUT\HEAD\PATCH\OPTIONS mode
	server.router.Handle(routeMethod, path, server.wrapRouterHandle(handle, false))
	return
}

// ServerFile is a shortcut for router.ServeFiles(path, filepath)
// simple demo:server.ServerFile("/src/*filepath", "/var/www")
func (server *HttpServer) ServerFile(urlpath string, filepath string) {
	server.router.ServeFiles(urlpath, http.Dir(filepath))
}

// WebSocket is a shortcut for websocket.Handler
func (server *HttpServer) WebSocket(path string, handle HttpHandle) {
	server.RegisterRoute(RouteMethod_WebSocket, path, handle)
}

type LogJson struct {
	RequestUrl string
	HttpHeader string
	HttpBody   string
}

//wrap HttpHandle to httprouter.Handle
func (server *HttpServer) wrapRouterHandle(handle HttpHandle, isHijack bool) httprouter.Handle {
	return func(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
		//get from pool
		res := server.pool.response.Get().(*Response)
		res.Reset(w)
		httpCtx := server.pool.context.Get().(*HttpContext)
		httpCtx.Reset(res, r, params)

		//增加状态计数
		GlobalState.AddRequestCount(1)

		//hijack处理
		if isHijack {
			_, hijack_err := httpCtx.Hijack()
			if hijack_err != nil {
				//输出内容
				httpCtx.Response.WriteHeader(http.StatusInternalServerError)
				httpCtx.Response.Header().Set(HeaderContentType, CharsetUTF8)
				httpCtx.WriteString(hijack_err.Error())
			}
		}

		startTime := time.Now()
		defer func() {
			var errmsg string
			if err := recover(); err != nil {
				errmsg = exception.CatchError("httpserver::RouterHandle", LogTarget_HttpServer, err)

				//默认异常处理
				if server.dotweb.ExceptionHandler != nil {
					server.dotweb.ExceptionHandler(httpCtx, err)
				}

				//记录访问日志
				headinfo := fmt.Sprintln(httpCtx.Response.Header)
				logJson := LogJson{
					RequestUrl: httpCtx.Request.RequestURI,
					HttpHeader: headinfo,
					HttpBody:   errmsg,
				}
				logString := jsonutil.GetJsonString(logJson)
				logger.Log(logString, LogTarget_HttpServer, LogLevel_Error)

				//增加错误计数
				GlobalState.AddErrorCount(1)
			}
			timetaken := int64(time.Now().Sub(startTime) / time.Millisecond)
			//HttpServer Logging
			logger.Log(httpCtx.Url()+" "+logString(httpCtx, timetaken), LogTarget_HttpRequest, LogLevel_Debug)

			// Return to pool
			server.pool.response.Put(res)
			server.pool.context.Put(httpCtx)
		}()

		//处理前置Module集合
		for _, module := range server.dotweb.Modules {
			if module.OnBeginRequest != nil {
				module.OnBeginRequest(httpCtx)
				//if current module set HttpContext.End,stop this request
				if httpCtx.isEnd {
					return
				}
			}
		}

		//处理用户handle
		handle(httpCtx)
		if httpCtx.isEnd {
			//if current module set HttpContext.End,stop this request
			return
		}

		//处理后置Module集合
		for _, module := range server.dotweb.Modules {
			if module.OnEndRequest != nil {
				module.OnEndRequest(httpCtx)
				if httpCtx.isEnd {
					//if current module set HttpContext.End,stop this request
					return
				}
			}
		}

	}
}

//wrap HttpHandle to websocket.Handle
func (server *HttpServer) wrapWebSocketHandle(handle HttpHandle) websocket.Handler {
	return func(ws *websocket.Conn) {
		//get from pool
		httpCtx := server.pool.context.Get().(*HttpContext)
		httpCtx.Reset(nil, ws.Request(), nil)
		httpCtx.WebSocket = &WebSocket{
			Conn: ws,
		}
		httpCtx.IsWebSocket = true

		startTime := time.Now()
		defer func() {
			var errmsg string
			if err := recover(); err != nil {
				errmsg = exception.CatchError("httpserver::WebsocketHandle", LogTarget_HttpServer, err)

				//记录访问日志
				headinfo := fmt.Sprintln(httpCtx.WebSocket.Conn.Request().Header)
				logJson := LogJson{
					RequestUrl: httpCtx.WebSocket.Conn.Request().RequestURI,
					HttpHeader: headinfo,
					HttpBody:   errmsg,
				}
				logString := jsonutil.GetJsonString(logJson)
				logger.Log(logString, LogTarget_HttpServer, LogLevel_Error)

				//增加错误计数
				GlobalState.AddErrorCount(1)
			}
			timetaken := int64(time.Now().Sub(startTime) / time.Millisecond)
			//HttpServer Logging
			logger.Log(httpCtx.Url()+" "+logString(httpCtx, timetaken), LogTarget_HttpRequest, LogLevel_Debug)

			// Return to pool
			server.pool.context.Put(httpCtx)
		}()

		handle(httpCtx)

		//增加状态计数
		GlobalState.AddRequestCount(1)
	}
}

//get default log string
func logString(ctx *HttpContext, timetaken int64) string {

	var reqbytelen, resbytelen, method, proto, status, userip string
	if !ctx.IsWebSocket {
		reqbytelen = convert.Int642String(ctx.Request.ContentLength)
		resbytelen = convert.Int642String(ctx.Response.Size)
		method = ctx.Request.Method
		proto = ctx.Request.Proto
		status = convert.Int2String(ctx.Response.Status)
		userip = ctx.RemoteIP()
	} else {
		reqbytelen = convert.Int642String(ctx.Request.ContentLength)
		resbytelen = "0"
		method = ctx.Request.Method
		proto = ctx.Request.Proto
		status = "0"
		userip = ctx.RemoteIP()
	}

	log := method + " "
	log += userip + " "
	log += proto + " "
	log += status + " "
	log += reqbytelen + " "
	log += resbytelen + " "
	log += convert.Int642String(timetaken)

	return log
}

//check request is the websocket request
func checkIsWebSocketRequest(req *http.Request) bool {
	if req.Header.Get("Connection") == "Upgrade" {
		return true
	}
	return false
}

//check request is startwith /debug/
func checkIsDebugRequest(req *http.Request) bool {
	if strings.Index(req.RequestURI, "/debug/") == 0 {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// DNS server implementation.

package dns

import (
	"io"
	"net"
	"radix"
	"time"
)

type Handler interface {
	ServeDNS(w ResponseWriter, r *Msg)
}

// A ResponseWriter interface is used by an DNS handler to
// construct an DNS response.
type ResponseWriter interface {
	// RemoteAddr returns the net.Addr of the client that sent the current request.
	RemoteAddr() net.Addr
	// TsigSttus returns the status of the Tsig (TsigNone, TsigVerified or TsigBad).
	TsigStatus() error
	// Write writes a reply back to the client.
	Write(*Msg) error
}

type conn struct {
	remoteAddr net.Addr          // address of the client
	handler    Handler           // request handler
	request    []byte            // bytes read
	_UDP       *net.UDPConn      // i/o connection if UDP was used
	_TCP       *net.TCPConn      // i/o connection if TCP was used
	hijacked   bool              // connection has been hijacked by hander TODO(mg)
	tsigSecret map[string]string // the tsig secrets
}

type response struct {
	conn           *conn
	req            *Msg
	tsigStatus     error
	tsigTimersOnly bool
	tsigRequestMAC string
}

// ServeMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name. ServeMux is DNSSEC aware, meaning
// that queries for the DS record are redirected to the parent zone (if that
// is also registered).
type ServeMux struct {
	m *radix.Radix
}

// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux { return &ServeMux{m: radix.New()} }

// DefaultServeMux is the default ServeMux used by Serve.
var DefaultServeMux = NewServeMux()

// The HandlerFunc type is an adapter to allow the use of
// ordinary functions as DNS handlers.  If f is a function
// with the appropriate signature, HandlerFunc(f) is a
// Handler object that calls f.
type HandlerFunc func(ResponseWriter, *Msg)

// ServerDNS calls f(w, r)
func (f HandlerFunc) ServeDNS(w ResponseWriter, r *Msg) {
	f(w, r)
}

// Failed is a helper handler that returns an answer with
// RCODE = servfail for every request.
func Failed(w ResponseWriter, r *Msg) {
	m := new(Msg)
	m.SetRcode(r, RcodeServerFailure)
	w.Write(m)
}

// FailedHandler returns HandlerFunc with Failed.
func FailedHandler() Handler { return HandlerFunc(Failed) }

// Start a server on addresss and network speficied. Invoke handler
// for any incoming queries.
func ListenAndServe(addr string, network string, handler Handler) error {
	server := &Server{Addr: addr, Net: network, Handler: handler}
	return server.ListenAndServe()
}

// Start a server on addresss and network speficied. Use the tsig
// secrets for Tsig validation.
// Invoke handler for any incoming queries.
func ListenAndServeTsig(addr string, network string, handler Handler, tsig map[string]string) error {
	server := &Server{Addr: addr, Net: network, Handler: handler, TsigSecret: tsig}
	return server.ListenAndServe()
}

func (mux *ServeMux) match(zone string, t uint16) Handler {
	// Exact match
	zone = toRadixName(zone)
	if h := mux.m.Find(zone); h != nil && h.Value != nil {
		// If we got queried for a DS record, we must see if we
		// if we also serve the parent. We then redirect it.
		if t == TypeDS {
			if d := h.Up(); d != nil {
				return d.Value.(Handler)
			}
		}
		return h.Value.(Handler)
	}
	// Best matching one.
	if h := mux.m.Predecessor(zone); h != nil && h.Value != nil {
		return h.Value.(Handler)
	}
	return nil
}

// Handle adds a handler to the ServeMux for pattern.
func (mux *ServeMux) Handle(pattern string, handler Handler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	mux.m.Insert(toRadixName(Fqdn(pattern)), handler)
}

// Handle adds a handler to the ServeMux for pattern.
func (mux *ServeMux) HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	mux.Handle(pattern, HandlerFunc(handler))
}

// HandleRemove deregistrars the handler specific for pattern from the ServeMux.
func (mux *ServeMux) HandleRemove(pattern string) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	// if its there, its gone
	mux.m.Remove(toRadixName(Fqdn(pattern)))
}

// ServeDNS dispatches the request to the handler whose
// pattern most closely matches the request message. For DS queries
// a parent zone is sought.
// If no handler is found a standard SERVFAIL message is returned
func (mux *ServeMux) ServeDNS(w ResponseWriter, request *Msg) {
	h := mux.match(request.Question[0].Name, request.Question[0].Qtype)
	if h == nil {
		h = FailedHandler()
	}
	h.ServeDNS(w, request)
}

// Handle registers the handler with the given pattern
// in the DefaultServeMux. The documentation for
// ServeMux explains how patters are matched.
func Handle(pattern string, handler Handler) { DefaultServeMux.Handle(pattern, handler) }

// HandleRemove deregisters the handle with the given pattern
// in the DefaultServeMux.
func HandleRemove(pattern string) { DefaultServeMux.HandleRemove(pattern) }

// HandleFunc registers the handler function with te given pattern
// in the DefaultServeMux.
func HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	DefaultServeMux.HandleFunc(pattern, handler)
}

// A Server defines parameters for running an DNS server.
type Server struct {
	Addr         string            // address to listen on, ":dns" if empty
	Net          string            // if "tcp" it will invoke a TCP listener, otherwise an UDP one
	Handler      Handler           // handler to invoke, dns.DefaultServeMux if nil
	UDPSize      int               // default buffer to use to read incoming UDP messages
	ReadTimeout  time.Duration     // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout time.Duration     // the net.Conn.SetWriteTimeout value for new connections
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>
}

// ListenAndServe starts a nameserver on the configured address in *Server.
func (srv *Server) ListenAndServe() error {
	addr := srv.Addr
	if addr == "" {
		addr = ":domain"
	}
	switch srv.Net {
	case "tcp", "tcp4", "tcp6":
		a, e := net.ResolveTCPAddr(srv.Net, addr)
		if e != nil {
			return e
		}
		l, e := net.ListenTCP(srv.Net, a)
		if e != nil {
			return e
		}
		return srv.ServeTCP(l)
	case "udp", "udp4", "udp6":
		a, e := net.ResolveUDPAddr(srv.Net, addr)
		if e != nil {
			return e
		}
		l, e := net.ListenUDP(srv.Net, a)
		if e != nil {
			return e
		}
		return srv.ServeUDP(l)
	}
	return &Error{Err: "bad network"}
}

// ServeTCP starts a TCP listener for the server.
// Each request is handled in a seperate goroutine.
// with the Handler set in ....
func (srv *Server) ServeTCP(l *net.TCPListener) error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
forever:
	for {
		rw, e := l.AcceptTCP()
		if e != nil {
			return e
		}
		if srv.ReadTimeout != 0 {
			rw.SetReadDeadline(time.Now().Add(srv.ReadTimeout))
		}
		if srv.WriteTimeout != 0 {
			rw.SetWriteDeadline(time.Now().Add(srv.WriteTimeout))
		}
		l := make([]byte, 2)
		n, err := rw.Read(l)
		if err != nil || n != 2 {
			continue
		}
		length, _ := unpackUint16(l, 0)
		if length == 0 {
			continue
		}
		m := make([]byte, int(length))
		n, err = rw.Read(m[:int(length)])
		if err != nil || n == 0 {
			continue
		}
		i := n
		for i < int(length) {
			j, err := rw.Read(m[i:int(length)])
			if err != nil {
				continue forever
			}
			i += j
		}
		n = i
		d, err := newConn(rw, nil, rw.RemoteAddr(), m, handler, srv.TsigSecret)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

// ServeUDP starts a UDP listener for the server.
// Each request is handled in a seperate goroutine,
// with the Handler set in ....
func (srv *Server) ServeUDP(l *net.UDPConn) error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
	if srv.UDPSize == 0 {
		srv.UDPSize = UDPMsgSize
	}
	for {
		m := make([]byte, srv.UDPSize)
		n, a, e := l.ReadFromUDP(m)
		if e != nil || n == 0 {
			return e
		}
		m = m[:n]

		if srv.ReadTimeout != 0 {
			l.SetReadDeadline(time.Now().Add(srv.ReadTimeout))
		}
		if srv.WriteTimeout != 0 {
			l.SetWriteDeadline(time.Now().Add(srv.WriteTimeout))
		}
		d, err := newConn(nil, l, a, m, handler, srv.TsigSecret)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

func newConn(t *net.TCPConn, u *net.UDPConn, a net.Addr, buf []byte, handler Handler, tsig map[string]string) (*conn, error) {
	c := new(conn)
	c.handler = handler
	c._TCP = t
	c._UDP = u
	c.remoteAddr = a
	c.request = buf
	c.tsigSecret = tsig
	return c, nil
}

// Close the connection.
func (c *conn) close() {
	switch {
	case c._UDP != nil:
		c._UDP.Close()
		c._UDP = nil
	case c._TCP != nil:
		c._TCP.Close()
		c._TCP = nil
	}
}

// Serve a new connection.
func (c *conn) serve() {
	for {
		// Request has been read in ServeUDP or ServeTCP
		w := new(response)
		w.conn = c
		req := new(Msg)
		if !req.Unpack(c.request) {
			// Send a format error back
			x := new(Msg)
			x.SetRcodeFormatError(req)
			w.Write(x)
			break
		}

		w.tsigStatus = nil
		if req.IsTsig() {
			secret := req.Extra[len(req.Extra)-1].(*RR_TSIG).Hdr.Name
			if _, ok := w.conn.tsigSecret[secret]; !ok {
				w.tsigStatus = ErrKeyAlg
			}
			w.tsigStatus = TsigVerify(c.request, w.conn.tsigSecret[secret], "", false)
			w.tsigTimersOnly = false // Will this ever be true?
			w.tsigRequestMAC = req.Extra[len(req.Extra)-1].(*RR_TSIG).MAC
		}
		w.req = req
		c.handler.ServeDNS(w, w.req) // this does the writing back to the client
		if c.hijacked {
			return
		}
		break // TODO(mg) Why is this a loop anyway?
	}
	if c._TCP != nil {
		c.close() // Listen and Serve is closed then
	}
}

func (w *response) Write(m *Msg) (err error) {
	var (
		data []byte
		ok   bool
	)
	if m == nil {
		return &Error{Err: "nil message"}
	}
	if m.IsTsig() {
		data, w.tsigRequestMAC, err = TsigGenerate(m, w.conn.tsigSecret[m.Extra[len(m.Extra)-1].(*RR_TSIG).Hdr.Name], w.tsigRequestMAC, w.tsigTimersOnly)
		if err != nil {
			return err
		}
	} else {
		data, ok = m.Pack()
		if !ok {
			return ErrPack
		}
	}
	switch {
	case w.conn._UDP != nil:
		_, err := w.conn._UDP.WriteTo(data, w.conn.remoteAddr)
		if err != nil {
			return err
		}
	case w.conn._TCP != nil:
		if len(data) > MaxMsgSize {
			return ErrBuf
		}
		l := make([]byte, 2)
		l[0], l[1] = packUint16(uint16(len(data)))
		n, err := w.conn._TCP.Write(l)
		if err != nil {
			return err
		}
		if n != 2 {
			return io.ErrShortWrite
		}
		n, err = w.conn._TCP.Write(data)
		if err != nil {
			return err
		}
		i := n
		if i < len(data) {
			j, err := w.conn._TCP.Write(data[i:len(data)])
			if err != nil {
				return err
			}
			i += j
		}
		n = i
	}
	return nil
}

// RemoteAddr implements the ResponseWriter.RemoteAddr method
func (w *response) RemoteAddr() net.Addr { return w.conn.remoteAddr }

// TsigStatus implements the ResponseWriter.TsigStatus method
func (w *response) TsigStatus() error { return w.tsigStatus }
<commit_msg>Add writebuf<commit_after>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// DNS server implementation.

package dns

import (
	"io"
	"net"
	"radix"
	"time"
)

type Handler interface {
	ServeDNS(w ResponseWriter, r *Msg)
}

// A ResponseWriter interface is used by an DNS handler to
// construct an DNS response.
type ResponseWriter interface {
	// RemoteAddr returns the net.Addr of the client that sent the current request.
	RemoteAddr() net.Addr
	// TsigSttus returns the status of the Tsig (TsigNone, TsigVerified or TsigBad).
	TsigStatus() error
	// Write writes a reply back to the client.
	Write(*Msg) error
	// WriteBuf writes a raw buffer back to the client.
	WriteBuf([]byte) error
}

type conn struct {
	remoteAddr net.Addr          // address of the client
	handler    Handler           // request handler
	request    []byte            // bytes read
	_UDP       *net.UDPConn      // i/o connection if UDP was used
	_TCP       *net.TCPConn      // i/o connection if TCP was used
	hijacked   bool              // connection has been hijacked by hander TODO(mg)
	tsigSecret map[string]string // the tsig secrets
}

type response struct {
	conn           *conn
	req            *Msg
	tsigStatus     error
	tsigTimersOnly bool
	tsigRequestMAC string
}

// ServeMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name. ServeMux is DNSSEC aware, meaning
// that queries for the DS record are redirected to the parent zone (if that
// is also registered).
type ServeMux struct {
	m *radix.Radix
}

// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux { return &ServeMux{m: radix.New()} }

// DefaultServeMux is the default ServeMux used by Serve.
var DefaultServeMux = NewServeMux()

// The HandlerFunc type is an adapter to allow the use of
// ordinary functions as DNS handlers.  If f is a function
// with the appropriate signature, HandlerFunc(f) is a
// Handler object that calls f.
type HandlerFunc func(ResponseWriter, *Msg)

// ServerDNS calls f(w, r)
func (f HandlerFunc) ServeDNS(w ResponseWriter, r *Msg) {
	f(w, r)
}

// Failed is a helper handler that returns an answer with
// RCODE = servfail for every request.
func Failed(w ResponseWriter, r *Msg) {
	m := new(Msg)
	m.SetRcode(r, RcodeServerFailure)
	w.Write(m)
}

// FailedHandler returns HandlerFunc with Failed.
func FailedHandler() Handler { return HandlerFunc(Failed) }

// Start a server on addresss and network speficied. Invoke handler
// for any incoming queries.
func ListenAndServe(addr string, network string, handler Handler) error {
	server := &Server{Addr: addr, Net: network, Handler: handler}
	return server.ListenAndServe()
}

// Start a server on addresss and network speficied. Use the tsig
// secrets for Tsig validation.
// Invoke handler for any incoming queries.
func ListenAndServeTsig(addr string, network string, handler Handler, tsig map[string]string) error {
	server := &Server{Addr: addr, Net: network, Handler: handler, TsigSecret: tsig}
	return server.ListenAndServe()
}

func (mux *ServeMux) match(zone string, t uint16) Handler {
	// Exact match
	zone = toRadixName(zone)
	if h := mux.m.Find(zone); h != nil && h.Value != nil {
		// If we got queried for a DS record, we must see if we
		// if we also serve the parent. We then redirect it.
		if t == TypeDS {
			if d := h.Up(); d != nil {
				return d.Value.(Handler)
			}
		}
		return h.Value.(Handler)
	}
	// Best matching one.
	if h := mux.m.Predecessor(zone); h != nil && h.Value != nil {
		return h.Value.(Handler)
	}
	return nil
}

// Handle adds a handler to the ServeMux for pattern.
func (mux *ServeMux) Handle(pattern string, handler Handler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	mux.m.Insert(toRadixName(Fqdn(pattern)), handler)
}

// Handle adds a handler to the ServeMux for pattern.
func (mux *ServeMux) HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	mux.Handle(pattern, HandlerFunc(handler))
}

// HandleRemove deregistrars the handler specific for pattern from the ServeMux.
func (mux *ServeMux) HandleRemove(pattern string) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	// if its there, its gone
	mux.m.Remove(toRadixName(Fqdn(pattern)))
}

// ServeDNS dispatches the request to the handler whose
// pattern most closely matches the request message. For DS queries
// a parent zone is sought.
// If no handler is found a standard SERVFAIL message is returned
func (mux *ServeMux) ServeDNS(w ResponseWriter, request *Msg) {
	h := mux.match(request.Question[0].Name, request.Question[0].Qtype)
	if h == nil {
		h = FailedHandler()
	}
	h.ServeDNS(w, request)
}

// Handle registers the handler with the given pattern
// in the DefaultServeMux. The documentation for
// ServeMux explains how patters are matched.
func Handle(pattern string, handler Handler) { DefaultServeMux.Handle(pattern, handler) }

// HandleRemove deregisters the handle with the given pattern
// in the DefaultServeMux.
func HandleRemove(pattern string) { DefaultServeMux.HandleRemove(pattern) }

// HandleFunc registers the handler function with te given pattern
// in the DefaultServeMux.
func HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	DefaultServeMux.HandleFunc(pattern, handler)
}

// A Server defines parameters for running an DNS server.
type Server struct {
	Addr         string            // address to listen on, ":dns" if empty
	Net          string            // if "tcp" it will invoke a TCP listener, otherwise an UDP one
	Handler      Handler           // handler to invoke, dns.DefaultServeMux if nil
	UDPSize      int               // default buffer to use to read incoming UDP messages
	ReadTimeout  time.Duration     // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout time.Duration     // the net.Conn.SetWriteTimeout value for new connections
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>
}

// ListenAndServe starts a nameserver on the configured address in *Server.
func (srv *Server) ListenAndServe() error {
	addr := srv.Addr
	if addr == "" {
		addr = ":domain"
	}
	switch srv.Net {
	case "tcp", "tcp4", "tcp6":
		a, e := net.ResolveTCPAddr(srv.Net, addr)
		if e != nil {
			return e
		}
		l, e := net.ListenTCP(srv.Net, a)
		if e != nil {
			return e
		}
		return srv.ServeTCP(l)
	case "udp", "udp4", "udp6":
		a, e := net.ResolveUDPAddr(srv.Net, addr)
		if e != nil {
			return e
		}
		l, e := net.ListenUDP(srv.Net, a)
		if e != nil {
			return e
		}
		return srv.ServeUDP(l)
	}
	return &Error{Err: "bad network"}
}

// ServeTCP starts a TCP listener for the server.
// Each request is handled in a seperate goroutine.
// with the Handler set in ....
func (srv *Server) ServeTCP(l *net.TCPListener) error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
forever:
	for {
		rw, e := l.AcceptTCP()
		if e != nil {
			return e
		}
		if srv.ReadTimeout != 0 {
			rw.SetReadDeadline(time.Now().Add(srv.ReadTimeout))
		}
		if srv.WriteTimeout != 0 {
			rw.SetWriteDeadline(time.Now().Add(srv.WriteTimeout))
		}
		l := make([]byte, 2)
		n, err := rw.Read(l)
		if err != nil || n != 2 {
			continue
		}
		length, _ := unpackUint16(l, 0)
		if length == 0 {
			continue
		}
		m := make([]byte, int(length))
		n, err = rw.Read(m[:int(length)])
		if err != nil || n == 0 {
			continue
		}
		i := n
		for i < int(length) {
			j, err := rw.Read(m[i:int(length)])
			if err != nil {
				continue forever
			}
			i += j
		}
		n = i
		d, err := newConn(rw, nil, rw.RemoteAddr(), m, handler, srv.TsigSecret)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

// ServeUDP starts a UDP listener for the server.
// Each request is handled in a seperate goroutine,
// with the Handler set in ....
func (srv *Server) ServeUDP(l *net.UDPConn) error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
	if srv.UDPSize == 0 {
		srv.UDPSize = UDPMsgSize
	}
	for {
		m := make([]byte, srv.UDPSize)
		n, a, e := l.ReadFromUDP(m)
		if e != nil || n == 0 {
			return e
		}
		m = m[:n]

		if srv.ReadTimeout != 0 {
			l.SetReadDeadline(time.Now().Add(srv.ReadTimeout))
		}
		if srv.WriteTimeout != 0 {
			l.SetWriteDeadline(time.Now().Add(srv.WriteTimeout))
		}
		d, err := newConn(nil, l, a, m, handler, srv.TsigSecret)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

func newConn(t *net.TCPConn, u *net.UDPConn, a net.Addr, buf []byte, handler Handler, tsig map[string]string) (*conn, error) {
	c := new(conn)
	c.handler = handler
	c._TCP = t
	c._UDP = u
	c.remoteAddr = a
	c.request = buf
	c.tsigSecret = tsig
	return c, nil
}

// Close the connection.
func (c *conn) close() {
	switch {
	case c._UDP != nil:
		c._UDP.Close()
		c._UDP = nil
	case c._TCP != nil:
		c._TCP.Close()
		c._TCP = nil
	}
}

// Serve a new connection.
func (c *conn) serve() {
	for {
		// Request has been read in ServeUDP or ServeTCP
		w := new(response)
		w.conn = c
		req := new(Msg)
		if !req.Unpack(c.request) {
			// Send a format error back
			x := new(Msg)
			x.SetRcodeFormatError(req)
			w.Write(x)
			break
		}

		w.tsigStatus = nil
		if req.IsTsig() {
			secret := req.Extra[len(req.Extra)-1].(*RR_TSIG).Hdr.Name
			if _, ok := w.conn.tsigSecret[secret]; !ok {
				w.tsigStatus = ErrKeyAlg
			}
			w.tsigStatus = TsigVerify(c.request, w.conn.tsigSecret[secret], "", false)
			w.tsigTimersOnly = false // Will this ever be true?
			w.tsigRequestMAC = req.Extra[len(req.Extra)-1].(*RR_TSIG).MAC
		}
		w.req = req
		c.handler.ServeDNS(w, w.req) // this does the writing back to the client
		if c.hijacked {
			return
		}
		break // TODO(mg) Why is this a loop anyway?
	}
	if c._TCP != nil {
		c.close() // Listen and Serve is closed then
	}
}

func (w *response) Write(m *Msg) (err error) {
	var (
		data []byte
		ok   bool
	)
	if m == nil {
		return &Error{Err: "nil message"}
	}
	if m.IsTsig() {
		data, w.tsigRequestMAC, err = TsigGenerate(m, w.conn.tsigSecret[m.Extra[len(m.Extra)-1].(*RR_TSIG).Hdr.Name], w.tsigRequestMAC, w.tsigTimersOnly)
		if err != nil {
			return err
		}
	} else {
		data, ok = m.Pack()
		if !ok {
			return ErrPack
		}
	}
	switch {
	case w.conn._UDP != nil:
		_, err := w.conn._UDP.WriteTo(data, w.conn.remoteAddr)
		if err != nil {
			return err
		}
	case w.conn._TCP != nil:
		if len(data) > MaxMsgSize {
			return ErrBuf
		}
		l := make([]byte, 2)
		l[0], l[1] = packUint16(uint16(len(data)))
		n, err := w.conn._TCP.Write(l)
		if err != nil {
			return err
		}
		if n != 2 {
			return io.ErrShortWrite
		}
		n, err = w.conn._TCP.Write(data)
		if err != nil {
			return err
		}
		i := n
		if i < len(data) {
			j, err := w.conn._TCP.Write(data[i:len(data)])
			if err != nil {
				return err
			}
			i += j
		}
		n = i
	}
	return nil
}

func (w *response) WriteBuf(m []byte) (err error) {
	// TODO(mg): refacter as we duplicate code from above?
	if m == nil {
		return &Error{Err: "nil message"}
	}
	switch {
	case w.conn._UDP != nil:
		_, err := w.conn._UDP.WriteTo(m, w.conn.remoteAddr)
		if err != nil {
			return err
		}
	case w.conn._TCP != nil:
		if len(m) > MaxMsgSize {
			return ErrBuf
		}
		l := make([]byte, 2)
		l[0], l[1] = packUint16(uint16(len(m)))
		n, err := w.conn._TCP.Write(l)
		if err != nil {
			return err
		}
		if n != 2 {
			return io.ErrShortWrite
		}
		n, err = w.conn._TCP.Write(m)
		if err != nil {
			return err
		}
		i := n
		if i < len(m) {
			j, err := w.conn._TCP.Write(m[i:len(m)])
			if err != nil {
				return err
			}
			i += j
		}
		n = i
	}
	return nil
}

// RemoteAddr implements the ResponseWriter.RemoteAddr method
func (w *response) RemoteAddr() net.Addr { return w.conn.remoteAddr }

// TsigStatus implements the ResponseWriter.TsigStatus method
func (w *response) TsigStatus() error { return w.tsigStatus }
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

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

type Server struct {
	ID     string
	Config *Configuration
	Store  *Storage

	timeout time.Duration
}

func createServer(conf *Configuration, store *Storage) *Server {
	hash := md5.New()
	io.WriteString(hash, time.Now().String())
	id := string(hash.Sum(nil))

	timeout := time.Duration(conf.GetInt("connection_timeout"))
	return &Server{id, conf, store, timeout}
}

func (this *Server) initSocketListener() {
	Connect := func(ws *websocket.Conn) {
		defer func() {
			if DEBUG {
				log.Println("Socket Closed")
			}
		}()

		sock := newSocket(ws, nil, this, "")

		if DEBUG {
			log.Printf("Socket connected via %s\n", ws.RemoteAddr())
		}
		if err := sock.Authenticate(""); err != nil {
			if DEBUG {
				log.Printf("Error: %s\n", err.Error())
			}
			return
		}

		go sock.listenForMessages()
		go sock.listenForWrites()

		if this.timeout <= 0 { // if timeout is 0 then wait forever and return when socket is done.
			<-sock.done
			return
		}

		select {
		case <-time.After(this.timeout * time.Second):
			sock.Close()
			return
		case <-sock.done:
			return
		}
	}

	http.Handle("/socket", websocket.Handler(Connect))
}

func (this *Server) initLongPollListener() {
	LpConnect := func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if DEBUG {
				log.Println("Socket Closed")
			}
		}()

		sock := newSocket(nil, w, this, "")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Cache-Control", "private, no-store, no-cache, must-revalidate, post-check=0, pre-check=0")
		w.Header().Set("Connection", "keep-alive")
		//w.Header().Set("Content-Encoding", "gzip")
		w.WriteHeader(200)

		if DEBUG {
			log.Printf("Long poll connected via \n")
		}

		if err := sock.Authenticate(r.FormValue("user")); err != nil {
			if DEBUG {
				log.Printf("Error: %s\n", err.Error())
			}
			return
		}

		page := r.FormValue("page")
		if page != "" {
			if sock.Page != "" {
				this.Store.UnsetPage(sock) //remove old page if it exists
			}

			sock.Page = page
			this.Store.SetPage(sock)
		}

		command := r.FormValue("command")
		if command != "" {
			var cmd CommandMsg
			json.Unmarshal([]byte(command), &cmd)

			go cmd.FromSocket(sock)
		}

		go sock.listenForWrites()

		select {
		case <-time.After(30 * time.Second):
			sock.Close()
			return
		case <-sock.done:
			return
		}
	}

	http.HandleFunc("/lp", LpConnect)
}

func (this *Server) initAppListener() {
	rec := make(chan []string)
	consumer, err := this.Store.redis.Subscribe(rec, this.Config.Get("redis_message_channel"))
	if err != nil {
		log.Fatal("Couldn't subscribe to redis channel")
	}
	defer consumer.Quit()

	if DEBUG {
		log.Println("LISENING FOR REDIS MESSAGE")
	}
	var ms []string
	for {
		ms = <-rec

		var cmd CommandMsg
		json.Unmarshal([]byte(ms[2]), &cmd)
		go cmd.FromRedis(this)
	}
}

func (this *Server) initPingListener() {
	pingHandler := func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "OK")
	}

	http.HandleFunc("/ping", pingHandler)
}

func (this *Server) sendHeartbeats() {

	for {
		time.Sleep(20 * time.Second)

		clients := this.Store.Clients()

		for _, user := range clients {
			for _, sock := range user {

				if sock.isWebsocket() {
					websocket.Message.Send(sock.ws, "")
				}

			}
		}
	}
}
<commit_msg>make the redis listener a buffered channel<commit_after>package main

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

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

type Server struct {
	ID     string
	Config *Configuration
	Store  *Storage

	timeout time.Duration
}

func createServer(conf *Configuration, store *Storage) *Server {
	hash := md5.New()
	io.WriteString(hash, time.Now().String())
	id := string(hash.Sum(nil))

	timeout := time.Duration(conf.GetInt("connection_timeout"))
	return &Server{id, conf, store, timeout}
}

func (this *Server) initSocketListener() {
	Connect := func(ws *websocket.Conn) {
		defer func() {
			if DEBUG {
				log.Println("Socket Closed")
			}
		}()

		sock := newSocket(ws, nil, this, "")

		if DEBUG {
			log.Printf("Socket connected via %s\n", ws.RemoteAddr())
		}
		if err := sock.Authenticate(""); err != nil {
			if DEBUG {
				log.Printf("Error: %s\n", err.Error())
			}
			return
		}

		go sock.listenForMessages()
		go sock.listenForWrites()

		if this.timeout <= 0 { // if timeout is 0 then wait forever and return when socket is done.
			<-sock.done
			return
		}

		select {
		case <-time.After(this.timeout * time.Second):
			sock.Close()
			return
		case <-sock.done:
			return
		}
	}

	http.Handle("/socket", websocket.Handler(Connect))
}

func (this *Server) initLongPollListener() {
	LpConnect := func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if DEBUG {
				log.Println("Socket Closed")
			}
		}()

		sock := newSocket(nil, w, this, "")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Cache-Control", "private, no-store, no-cache, must-revalidate, post-check=0, pre-check=0")
		w.Header().Set("Connection", "keep-alive")
		//w.Header().Set("Content-Encoding", "gzip")
		w.WriteHeader(200)

		if DEBUG {
			log.Printf("Long poll connected via \n")
		}

		if err := sock.Authenticate(r.FormValue("user")); err != nil {
			if DEBUG {
				log.Printf("Error: %s\n", err.Error())
			}
			return
		}

		page := r.FormValue("page")
		if page != "" {
			if sock.Page != "" {
				this.Store.UnsetPage(sock) //remove old page if it exists
			}

			sock.Page = page
			this.Store.SetPage(sock)
		}

		command := r.FormValue("command")
		if command != "" {
			var cmd CommandMsg
			json.Unmarshal([]byte(command), &cmd)

			go cmd.FromSocket(sock)
		}

		go sock.listenForWrites()

		select {
		case <-time.After(30 * time.Second):
			sock.Close()
			return
		case <-sock.done:
			return
		}
	}

	http.HandleFunc("/lp", LpConnect)
}

func (this *Server) initAppListener() {
	rec := make(chan []string, 10000)
	consumer, err := this.Store.redis.Subscribe(rec, this.Config.Get("redis_message_channel"))
	if err != nil {
		log.Fatal("Couldn't subscribe to redis channel")
	}
	defer consumer.Quit()

	if DEBUG {
		log.Println("LISENING FOR REDIS MESSAGE")
	}
	var ms []string
	for {
		ms = <-rec

		var cmd CommandMsg
		json.Unmarshal([]byte(ms[2]), &cmd)
		go cmd.FromRedis(this)
	}
}

func (this *Server) initPingListener() {
	pingHandler := func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "OK")
	}

	http.HandleFunc("/ping", pingHandler)
}

func (this *Server) sendHeartbeats() {

	for {
		time.Sleep(20 * time.Second)

		clients := this.Store.Clients()

		for _, user := range clients {
			for _, sock := range user {

				if sock.isWebsocket() {
					websocket.Message.Send(sock.ws, "")
				}

			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"

	"github.com/byuoitav/authmiddleware"
	"github.com/byuoitav/av-api/base"
	"github.com/byuoitav/av-api/handlers"
	avapi "github.com/byuoitav/av-api/init"
	"github.com/byuoitav/event-router-microservice/eventinfrastructure"
	"github.com/byuoitav/hateoas"
	"github.com/jessemillar/health"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"github.com/xuther/go-message-router/publisher"
)

func main() {
	//First we need to check if we're a room.
	err := avapi.CheckRoomInitialization()
	if err != nil {
		base.Publish(eventinfrastructure.Event{Event: "Fail to run init script. Terminating."})

		log.Fatalf("Could not initialize room. Error: %v\n", err.Error())
	}

	err = hateoas.Load("https://raw.githubusercontent.com/byuoitav/av-api/master/swagger.json")
	if err != nil {
		base.Publish(eventinfrastructure.Event{Event: "Fail to run init script. Terminating."})

		log.Fatalf("Could not load Swagger file. Error: %s\b", err.Error())
	}

	base.Publisher, err = publisher.NewPublisher("7001", 1000, 10)
	if err != nil {
		log.Fatalf("Could not start published. Error: %v\n", err.Error())
	} else {
		log.Printf("Publisher started on port :7001")
	}

	port := ":8000"
	router := echo.New()
	router.Pre(middleware.RemoveTrailingSlash())
	router.Use(middleware.CORS())

	// Use the `secure` routing group to require authentication
	secure := router.Group("", echo.WrapMiddleware(authmiddleware.Authenticate))

	// GET requests
	router.GET("/", echo.WrapHandler(http.HandlerFunc(hateoas.RootResponse)))

	router.GET("/health", echo.WrapHandler(http.HandlerFunc(health.Check)))

	// router.Get("/buildings", handlers.GetAllBuildings, wso2jwt.ValidateJWT())
	secure.GET("/buildings/:building/rooms/:room", handlers.GetRoomByNameAndBuilding)

	// PUT requests
	secure.PUT("/buildings/:building/rooms/:room", handlers.SetRoomState)

	// room status
	secure.GET("/buildings/:building/rooms/:room/status", handlers.GetRoomStatus)

	server := http.Server{
		Addr:           port,
		MaxHeaderBytes: 1024 * 10,
	}

	router.StartServer(&server)
}
<commit_msg>Listening<commit_after>package main

import (
	"log"
	"net/http"

	"github.com/byuoitav/authmiddleware"
	"github.com/byuoitav/av-api/base"
	"github.com/byuoitav/av-api/handlers"
	avapi "github.com/byuoitav/av-api/init"
	"github.com/byuoitav/event-router-microservice/eventinfrastructure"
	"github.com/byuoitav/hateoas"
	"github.com/jessemillar/health"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"github.com/xuther/go-message-router/publisher"
)

func main() {
	//First we need to check if we're a room.
	err := avapi.CheckRoomInitialization()
	if err != nil {
		base.Publish(eventinfrastructure.Event{Event: "Fail to run init script. Terminating."})

		log.Fatalf("Could not initialize room. Error: %v\n", err.Error())
	}

	err = hateoas.Load("https://raw.githubusercontent.com/byuoitav/av-api/master/swagger.json")
	if err != nil {
		base.Publish(eventinfrastructure.Event{Event: "Fail to run init script. Terminating."})

		log.Fatalf("Could not load Swagger file. Error: %s\b", err.Error())
	}

	base.Publisher, err = publisher.NewPublisher("7001", 1000, 10)
	if err != nil {
		log.Fatalf("Could not start published. Error: %v\n", err.Error())
	}

	go func() {
		base.Publisher.Listen()
		if err != nil {
			log.Fatalf("Could not start published. Error: %v\n", err.Error())
		} else {
			log.Printf("Publisher started on port :7001")
		}
	}()
	port := ":8000"
	router := echo.New()
	router.Pre(middleware.RemoveTrailingSlash())
	router.Use(middleware.CORS())

	// Use the `secure` routing group to require authentication
	secure := router.Group("", echo.WrapMiddleware(authmiddleware.Authenticate))

	// GET requests
	router.GET("/", echo.WrapHandler(http.HandlerFunc(hateoas.RootResponse)))

	router.GET("/health", echo.WrapHandler(http.HandlerFunc(health.Check)))

	// router.Get("/buildings", handlers.GetAllBuildings, wso2jwt.ValidateJWT())
	secure.GET("/buildings/:building/rooms/:room", handlers.GetRoomByNameAndBuilding)

	// PUT requests
	secure.PUT("/buildings/:building/rooms/:room", handlers.SetRoomState)

	// room status
	secure.GET("/buildings/:building/rooms/:room/status", handlers.GetRoomStatus)

	server := http.Server{
		Addr:           port,
		MaxHeaderBytes: 1024 * 10,
	}

	router.StartServer(&server)
}
<|endoftext|>


<commit_before>package cistatus

import (
	"encoding/json"
	"log"
	"net/http"
	"strings"
	"time"

	"sync"

	"io/ioutil"

	"github.com/dgrijalva/jwt-go"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
)

type Server struct {
	Logger *log.Logger

	JWT struct {
		Algorithm string
		Secret    []byte
	}

	*http.ServeMux

	fetcher       Fetcher
	latestSummary Summary
	websocket     struct {
		upgrader    websocket.Upgrader
		mutex       sync.RWMutex
		connections map[*websocket.Conn]bool
		broadcast   chan Summary
	}
}

func NewServer(fetcher Fetcher, fetchInterval time.Duration) *Server {
	now := time.Now()

	s := &Server{
		fetcher: fetcher,
		// Initial status summary is "unkown"
		latestSummary: Summary{
			Color:       Unknown,
			LastUpdated: &now,
		},
		// Default to discarding logs
		Logger: log.New(ioutil.Discard, "", 0),
	}

	// Setup for websockets
	s.websocket.upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	s.websocket.connections = make(map[*websocket.Conn]bool)
	s.websocket.broadcast = make(chan Summary)
	go s.handleBroadcasts()

	// Create servemux with routes to http api
	s.ServeMux = http.NewServeMux()
	s.ServeMux.HandleFunc("/api", s.allProjects)
	s.ServeMux.HandleFunc("/api/watch", s.websocketSubscribeHandler)

	// Start fetching
	go s.fetchLoop(fetchInterval)

	return s
}

func (s *Server) allProjects(w http.ResponseWriter, r *http.Request) {
	latestSummary := s.latestSummary

	if !s.isAuthorized(r) {
		latestSummary.Projects = []Project{}
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Server", "https://github.com/tantalic/cistatus")
	w.Header().Set("X-Server-Version", Version)
	json.NewEncoder(w).Encode(latestSummary)
}

func (s *Server) jwtKey(token *jwt.Token) (interface{}, error) {
	if token.Method.Alg() != s.JWT.Algorithm {
		return nil, errors.Errorf("unexpected jwt algorithm: %v", token.Header["alg"])
	}

	return s.JWT.Secret, nil
}

func (s *Server) isAuthorized(r *http.Request) bool {
	// If the JWT algorithm OR secret is not set then all requests are authorized
	if s.JWT.Secret == nil || s.JWT.Algorithm == "" {
		return true
	}

	authorizationHeader := r.Header.Get("Authorization")
	if authorizationHeader == "" {
		// no header present == not authorized
		return false
	}

	headerParts := strings.Split(authorizationHeader, " ")
	if len(headerParts) != 2 || headerParts[0] != "bearer" {
		// malformed header == not authorized, token should be in format of "bearer XXXXXXXXXXXXXXXX"
		return false
	}

	token, err := jwt.Parse(headerParts[1], s.jwtKey)
	if err != nil {
		s.Logger.Printf("JWT error: %s\n", err)
		return false
	}

	return token.Valid
}

func (s *Server) websocketSubscribeHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := s.websocket.upgrader.Upgrade(w, r, nil)
	if err != nil {
		err := errors.Wrap(err, "unable to create websocket connection to subscribe to status updates")
		s.Logger.Println(err)
		return
	}

	s.addSubscriber(conn)
}

func (s *Server) fetchLoop(interval time.Duration) {
	ticker := time.Tick(interval)
	errCount := 0

	for {
		s.Logger.Println("Fetching CI server status")

		projects, err := s.fetcher.FetchStatus()
		if err != nil {
			s.Logger.Printf("Error fetching status: %s\n", err)
			errCount++
			if errCount >= 10 {

			}
			<-ticker
			continue
		}

		now := time.Now()
		s.latestSummary.Projects = projects
		s.latestSummary.LastUpdated = &now

		newColor := color(projects)
		if newColor != s.latestSummary.Color {
			s.latestSummary.Color = newColor
			s.websocket.broadcast <- s.latestSummary
		}

		s.Logger.Printf("Fetched %d projects\n", len(projects))
		<-ticker
	}
}

func (s *Server) handleBroadcasts() {
	for {
		// Grab the next message from the broadcast channel
		msg := <-s.websocket.broadcast
		msg.Projects = nil

		// Send it out to every client that is currently connected
		s.websocket.mutex.RLock()
		for conn := range s.websocket.connections {
			err := conn.WriteJSON(msg)
			if err != nil {
				s.removeSubscriber(conn)
			}
		}
		s.websocket.mutex.RUnlock()
	}
}

func (s *Server) addSubscriber(conn *websocket.Conn) {
	// Send current status immediately
	summary := s.latestSummary
	summary.Projects = nil
	conn.WriteJSON(summary)

	// Add to list of connections
	s.websocket.mutex.Lock()
	s.websocket.connections[conn] = true
	s.websocket.mutex.Unlock()
}

func (s *Server) removeSubscriber(conn *websocket.Conn) {
	s.websocket.mutex.Lock()
	delete(s.websocket.connections, conn)
	s.websocket.mutex.Unlock()

	conn.Close()
}

func color(projects []Project) Color {
	color := Green

	for _, project := range projects {
		for _, branch := range project.Branches {
			for _, status := range branch.Statuses {

				// If any status is failed return red immediately
				if status.Status == "failed" {
					return Red
				}

				// If any status is pending or running return yellow
				if status.Status == "pending" || status.Status == "running" {
					color = Yellow
				}

			}
		}
	}

	return color
}
<commit_msg>Fixes spelling error in comment<commit_after>package cistatus

import (
	"encoding/json"
	"log"
	"net/http"
	"strings"
	"time"

	"sync"

	"io/ioutil"

	"github.com/dgrijalva/jwt-go"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
)

type Server struct {
	Logger *log.Logger

	JWT struct {
		Algorithm string
		Secret    []byte
	}

	*http.ServeMux

	fetcher       Fetcher
	latestSummary Summary
	websocket     struct {
		upgrader    websocket.Upgrader
		mutex       sync.RWMutex
		connections map[*websocket.Conn]bool
		broadcast   chan Summary
	}
}

func NewServer(fetcher Fetcher, fetchInterval time.Duration) *Server {
	now := time.Now()

	s := &Server{
		fetcher: fetcher,
		// Initial status summary is "unknown"
		latestSummary: Summary{
			Color:       Unknown,
			LastUpdated: &now,
		},
		// Default to discarding logs
		Logger: log.New(ioutil.Discard, "", 0),
	}

	// Setup for websockets
	s.websocket.upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	s.websocket.connections = make(map[*websocket.Conn]bool)
	s.websocket.broadcast = make(chan Summary)
	go s.handleBroadcasts()

	// Create servemux with routes to http api
	s.ServeMux = http.NewServeMux()
	s.ServeMux.HandleFunc("/api", s.allProjects)
	s.ServeMux.HandleFunc("/api/watch", s.websocketSubscribeHandler)

	// Start fetching
	go s.fetchLoop(fetchInterval)

	return s
}

func (s *Server) allProjects(w http.ResponseWriter, r *http.Request) {
	latestSummary := s.latestSummary

	if !s.isAuthorized(r) {
		latestSummary.Projects = []Project{}
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Server", "https://github.com/tantalic/cistatus")
	w.Header().Set("X-Server-Version", Version)
	json.NewEncoder(w).Encode(latestSummary)
}

func (s *Server) jwtKey(token *jwt.Token) (interface{}, error) {
	if token.Method.Alg() != s.JWT.Algorithm {
		return nil, errors.Errorf("unexpected jwt algorithm: %v", token.Header["alg"])
	}

	return s.JWT.Secret, nil
}

func (s *Server) isAuthorized(r *http.Request) bool {
	// If the JWT algorithm OR secret is not set then all requests are authorized
	if s.JWT.Secret == nil || s.JWT.Algorithm == "" {
		return true
	}

	authorizationHeader := r.Header.Get("Authorization")
	if authorizationHeader == "" {
		// no header present == not authorized
		return false
	}

	headerParts := strings.Split(authorizationHeader, " ")
	if len(headerParts) != 2 || headerParts[0] != "bearer" {
		// malformed header == not authorized, token should be in format of "bearer XXXXXXXXXXXXXXXX"
		return false
	}

	token, err := jwt.Parse(headerParts[1], s.jwtKey)
	if err != nil {
		s.Logger.Printf("JWT error: %s\n", err)
		return false
	}

	return token.Valid
}

func (s *Server) websocketSubscribeHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := s.websocket.upgrader.Upgrade(w, r, nil)
	if err != nil {
		err := errors.Wrap(err, "unable to create websocket connection to subscribe to status updates")
		s.Logger.Println(err)
		return
	}

	s.addSubscriber(conn)
}

func (s *Server) fetchLoop(interval time.Duration) {
	ticker := time.Tick(interval)
	errCount := 0

	for {
		s.Logger.Println("Fetching CI server status")

		projects, err := s.fetcher.FetchStatus()
		if err != nil {
			s.Logger.Printf("Error fetching status: %s\n", err)
			errCount++
			if errCount >= 10 {

			}
			<-ticker
			continue
		}

		now := time.Now()
		s.latestSummary.Projects = projects
		s.latestSummary.LastUpdated = &now

		newColor := color(projects)
		if newColor != s.latestSummary.Color {
			s.latestSummary.Color = newColor
			s.websocket.broadcast <- s.latestSummary
		}

		s.Logger.Printf("Fetched %d projects\n", len(projects))
		<-ticker
	}
}

func (s *Server) handleBroadcasts() {
	for {
		// Grab the next message from the broadcast channel
		msg := <-s.websocket.broadcast
		msg.Projects = nil

		// Send it out to every client that is currently connected
		s.websocket.mutex.RLock()
		for conn := range s.websocket.connections {
			err := conn.WriteJSON(msg)
			if err != nil {
				s.removeSubscriber(conn)
			}
		}
		s.websocket.mutex.RUnlock()
	}
}

func (s *Server) addSubscriber(conn *websocket.Conn) {
	// Send current status immediately
	summary := s.latestSummary
	summary.Projects = nil
	conn.WriteJSON(summary)

	// Add to list of connections
	s.websocket.mutex.Lock()
	s.websocket.connections[conn] = true
	s.websocket.mutex.Unlock()
}

func (s *Server) removeSubscriber(conn *websocket.Conn) {
	s.websocket.mutex.Lock()
	delete(s.websocket.connections, conn)
	s.websocket.mutex.Unlock()

	conn.Close()
}

func color(projects []Project) Color {
	color := Green

	for _, project := range projects {
		for _, branch := range project.Branches {
			for _, status := range branch.Statuses {

				// If any status is failed return red immediately
				if status.Status == "failed" {
					return Red
				}

				// If any status is pending or running return yellow
				if status.Status == "pending" || status.Status == "running" {
					color = Yellow
				}

			}
		}
	}

	return color
}
<|endoftext|>


<commit_before>// Copyright © 2009--2013 The Web.go Authors
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// this file is about the actual handling of a request: it comes in, what
// happens? routing determines which handler is responsible and that is then
// wrapped appropriately and invoked.

package web

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"path"
	"regexp"
	"runtime"
	"strings"
	"time"

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

type route struct {
	rex     *regexp.Regexp
	method  string
	handler parametrizedHandler
}

type ServerConfig struct {
	StaticDirs   []string
	Addr         string
	Port         int
	CookieSecret string
	RecoverPanic bool
	Cert         string
	Key          string
	ColorOutput  bool
}

type Server struct {
	Config ServerConfig
	routes []*route
	Logger *log.Logger
	Env    map[string]interface{}
	// Save the listener so it can be closed
	l net.Listener
	// Passed verbatim to every handler on every request
	User interface{}
	// All requests are passed through this wrapper if defined
	Wrappers []Wrapper
}

var mainServer = NewServer()

// Configuration of the shared server
var Config = &mainServer.Config
var exeFile string

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

// Queue response wrapper that is called after all other wrappers
func (s *Server) AddWrapper(wrap Wrapper) {
	s.Wrappers = append(s.Wrappers, wrap)
}

func (s *Server) SetLogger(logger *log.Logger) {
	s.Logger = logger
}

func (s *Server) addRoute(rawrex string, method string, handler interface{}) {
	rex, err := regexp.Compile(rawrex)
	if err != nil {
		s.Logger.Printf("Error in route regex %q: %v", rawrex, err)
		return
	}
	s.routes = append(s.routes, &route{
		rex:     rex,
		method:  method,
		handler: fixHandlerSignature(handler),
	})
}

// Calls function with recover block. The first return value is whatever the
// function returns if it didnt panic. The second is what was passed to panic()
// if it did.
func (s *Server) safelyCall(f func() error) (softerr error, harderr interface{}) {
	defer func() {
		if err := recover(); err != nil {
			if !s.Config.RecoverPanic {
				// go back to panic
				s.Logger.Printf("Panic: %v", err)
				panic(err)
			} else {
				harderr = err
				s.Logger.Println("Handler crashed with error: ", err)
				for i := 1; ; i += 1 {
					_, file, line, ok := runtime.Caller(i)
					if !ok {
						break
					}
					s.Logger.Println(file, line)
				}
			}
		}
	}()
	return f(), nil
}

// Determine if this route matches this request purely on the basis of the method
func matchRouteMethods(req *http.Request, route *route) bool {
	if req.Method == route.method {
		return true
	}
	if req.Method == "HEAD" && route.method == "GET" {
		return true
	}
	if req.Header.Get("Upgrade") == "websocket" && route.method == "WEBSOCKET" {
		return true
	}
	return false
}

// If this request matches this route return the group matches from the regular
// expression otherwise return an empty slice. note on success the return value
// includes the entire match as the first element.
func matchRoute(req *http.Request, route *route) []string {
	if !matchRouteMethods(req, route) {
		return nil
	}
	match := route.rex.FindStringSubmatch(req.URL.Path)
	if match == nil || len(match[0]) != len(req.URL.Path) {
		return nil
	}
	return match
}

func findMatchingRoute(req *http.Request, routes []*route) (*route, []string) {
	for _, route := range routes {
		if match := matchRoute(req, route); match != nil {
			return route, match
		}
	}
	return nil, nil
}

// Apply the handler to this context and try to handle errors where possible
func (s *Server) applyHandler(f SimpleHandler, ctx *Context) (err error) {
	softerr, harderr := s.safelyCall(func() error {
		return f(ctx)
	})
	if harderr != nil {
		//there was an error or panic while calling the handler
		ctx.Abort(500, "Server Error")
		return fmt.Errorf("%v", harderr)
	}
	if softerr != nil {
		s.Logger.Printf("Handler returned error: %v", softerr)
		if werr, ok := softerr.(WebError); ok {
			ctx.Abort(werr.Code, werr.Error())
		} else {
			// Non-web errors are not leaked to the outside
			ctx.Abort(500, "Server Error")
			err = softerr
		}
	} else {
		// flush the writer by ensuring at least one Write call takes place
		_, err = ctx.Write([]byte{})
	}
	if err == nil {
		err = ctx.Response.Close()
	}
	return
}

func dirExists(dir string) bool {
	d, e := os.Stat(dir)
	switch {
	case e != nil:
		return false
	case !d.IsDir():
		return false
	}

	return true
}

func fileExists(dir string) bool {
	info, err := os.Stat(dir)
	if err != nil {
		return false
	}

	return !info.IsDir()
}

// Default
func defaultStaticDir() string {
	root, _ := path.Split(exeFile)
	return path.Join(root, "static")
}

// If this request corresponds to a static file return its path
func (s *Server) findFile(req *http.Request) string {
	//try to serve static files
	staticDirs := s.Config.StaticDirs
	if len(staticDirs) == 0 {
		staticDirs = []string{defaultStaticDir()}
	}
	for _, staticDir := range staticDirs {
		staticFile := path.Join(staticDir, req.URL.Path)
		if fileExists(staticFile) && (req.Method == "GET" || req.Method == "HEAD") {
			return staticFile
		}
	}

	// Try to serve index.html || index.htm
	indexFilenames := []string{"index.html", "index.htm"}
	for _, staticDir := range staticDirs {
		for _, indexFilename := range indexFilenames {
			if indexPath := path.Join(path.Join(staticDir, req.URL.Path), indexFilename); fileExists(indexPath) {
				return indexPath
			}
		}
	}
	return ""
}

// Fully clothed request handler
func (s *Server) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	ctx := &Context{
		Request:  req,
		RawBody:  nil,
		Params:   map[string]string{},
		Server:   s,
		Response: &ResponseWriter{ResponseWriter: w, BodyWriter: w},
		User:     s.User,
	}

	//log the request
	if s.Config.ColorOutput {
		s.Logger.Printf("\033[32;1m%s %s\033[0m", req.Method, req.URL.Path)
	} else {
		s.Logger.Printf("%s %s", req.Method, req.URL.Path)
	}

	//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]
		}
		if s.Config.ColorOutput {
			s.Logger.Printf("\033[37;1mParams: %v\033[0m\n", ctx.Params)
		} else {
			s.Logger.Printf("Params: %v\n", ctx.Params)
		}

	}

	var simpleh SimpleHandler
	route, match := findMatchingRoute(req, s.routes)
	if route != nil {
		if route.method == "WEBSOCKET" {
			// Wrap websocket handler
			openh := func(ctx *Context, args ...string) (err error) {
				// yo dawg we heard you like wrapped functions
				websocket.Handler(func(ws *websocket.Conn) {
					ctx.WebsockConn = ws
					err = route.handler(ctx, args...)
				}).ServeHTTP(ctx.Response, req)
				return err
			}
			simpleh = closeHandler(openh, match[1:]...)
		} else {
			simpleh = closeHandler(route.handler, match[1:]...)
		}
	} else if path := s.findFile(req); path != "" {
		// no custom handler found but there is a file with this name
		simpleh = func(ctx *Context) error {
			http.ServeFile(ctx.Response, ctx.Request, path)
			return nil
		}
	} else {
		// hopeless, 404
		simpleh = func(ctx *Context) error {
			return WebError{404, "Page not found"}
		}
	}
	for _, wrap := range s.Wrappers {
		simpleh = wrapHandler(wrap, simpleh)
	}
	s.applyHandler(simpleh, ctx)
	return
}

func webTime(t time.Time) string {
	ftime := t.Format(time.RFC1123)
	if strings.HasSuffix(ftime, "UTC") {
		ftime = ftime[0:len(ftime)-3] + "GMT"
	}
	return ftime
}

func NewServer() *Server {
	conf := ServerConfig{
		RecoverPanic: true,
		Cert:         "",
		Key:          "",
		ColorOutput:  true,
	}
	s := &Server{
		Config: conf,
		Logger: log.New(os.Stdout, "", log.Ldate|log.Ltime),
		Env:    map[string]interface{}{},
	}
	// Set some default headers
	s.AddWrapper(func(h SimpleHandler, ctx *Context) error {
		ctx.Header().Set("Server", "web.go")
		tm := time.Now().UTC()
		ctx.Header().Set("Date", webTime(tm))
		return h(ctx)
	})
	// Compress response data when possible
	s.AddWrapper(CompressWrapper)
	return s
}

// Package wide proxy functions for global web server object

// Stop the global web server
func Close() {
	mainServer.Close()
}

// Set a logger to be used by the global web server
func SetLogger(logger *log.Logger) {
	mainServer.Logger = logger
}

// The global web server as an object implementing the http.Handler interface
func GetHTTPHandler() http.Handler {
	return mainServer
}

func init() {
	// find the location of the executable
	arg0 := path.Clean(os.Args[0])
	wd, _ := os.Getwd()
	if strings.HasPrefix(arg0, "/") {
		exeFile = arg0
	} else {
		// TODO For robustness, search each directory in $PATH
		exeFile = path.Join(wd, arg0)
	}
	registerDefaultMimetypes()
}
<commit_msg>Do not compress data by default<commit_after>// Copyright © 2009--2013 The Web.go Authors
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// this file is about the actual handling of a request: it comes in, what
// happens? routing determines which handler is responsible and that is then
// wrapped appropriately and invoked.

package web

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"path"
	"regexp"
	"runtime"
	"strings"
	"time"

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

type route struct {
	rex     *regexp.Regexp
	method  string
	handler parametrizedHandler
}

type ServerConfig struct {
	StaticDirs   []string
	Addr         string
	Port         int
	CookieSecret string
	RecoverPanic bool
	Cert         string
	Key          string
	ColorOutput  bool
}

type Server struct {
	Config ServerConfig
	routes []*route
	Logger *log.Logger
	Env    map[string]interface{}
	// Save the listener so it can be closed
	l net.Listener
	// Passed verbatim to every handler on every request
	User interface{}
	// All requests are passed through this wrapper if defined
	Wrappers []Wrapper
}

var mainServer = NewServer()

// Configuration of the shared server
var Config = &mainServer.Config
var exeFile string

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

// Queue response wrapper that is called after all other wrappers
func (s *Server) AddWrapper(wrap Wrapper) {
	s.Wrappers = append(s.Wrappers, wrap)
}

func (s *Server) SetLogger(logger *log.Logger) {
	s.Logger = logger
}

func (s *Server) addRoute(rawrex string, method string, handler interface{}) {
	rex, err := regexp.Compile(rawrex)
	if err != nil {
		s.Logger.Printf("Error in route regex %q: %v", rawrex, err)
		return
	}
	s.routes = append(s.routes, &route{
		rex:     rex,
		method:  method,
		handler: fixHandlerSignature(handler),
	})
}

// Calls function with recover block. The first return value is whatever the
// function returns if it didnt panic. The second is what was passed to panic()
// if it did.
func (s *Server) safelyCall(f func() error) (softerr error, harderr interface{}) {
	defer func() {
		if err := recover(); err != nil {
			if !s.Config.RecoverPanic {
				// go back to panic
				s.Logger.Printf("Panic: %v", err)
				panic(err)
			} else {
				harderr = err
				s.Logger.Println("Handler crashed with error: ", err)
				for i := 1; ; i += 1 {
					_, file, line, ok := runtime.Caller(i)
					if !ok {
						break
					}
					s.Logger.Println(file, line)
				}
			}
		}
	}()
	return f(), nil
}

// Determine if this route matches this request purely on the basis of the method
func matchRouteMethods(req *http.Request, route *route) bool {
	if req.Method == route.method {
		return true
	}
	if req.Method == "HEAD" && route.method == "GET" {
		return true
	}
	if req.Header.Get("Upgrade") == "websocket" && route.method == "WEBSOCKET" {
		return true
	}
	return false
}

// If this request matches this route return the group matches from the regular
// expression otherwise return an empty slice. note on success the return value
// includes the entire match as the first element.
func matchRoute(req *http.Request, route *route) []string {
	if !matchRouteMethods(req, route) {
		return nil
	}
	match := route.rex.FindStringSubmatch(req.URL.Path)
	if match == nil || len(match[0]) != len(req.URL.Path) {
		return nil
	}
	return match
}

func findMatchingRoute(req *http.Request, routes []*route) (*route, []string) {
	for _, route := range routes {
		if match := matchRoute(req, route); match != nil {
			return route, match
		}
	}
	return nil, nil
}

// Apply the handler to this context and try to handle errors where possible
func (s *Server) applyHandler(f SimpleHandler, ctx *Context) (err error) {
	softerr, harderr := s.safelyCall(func() error {
		return f(ctx)
	})
	if harderr != nil {
		//there was an error or panic while calling the handler
		ctx.Abort(500, "Server Error")
		return fmt.Errorf("%v", harderr)
	}
	if softerr != nil {
		s.Logger.Printf("Handler returned error: %v", softerr)
		if werr, ok := softerr.(WebError); ok {
			ctx.Abort(werr.Code, werr.Error())
		} else {
			// Non-web errors are not leaked to the outside
			ctx.Abort(500, "Server Error")
			err = softerr
		}
	} else {
		// flush the writer by ensuring at least one Write call takes place
		_, err = ctx.Write([]byte{})
	}
	if err == nil {
		err = ctx.Response.Close()
	}
	return
}

func dirExists(dir string) bool {
	d, e := os.Stat(dir)
	switch {
	case e != nil:
		return false
	case !d.IsDir():
		return false
	}

	return true
}

func fileExists(dir string) bool {
	info, err := os.Stat(dir)
	if err != nil {
		return false
	}

	return !info.IsDir()
}

// Default
func defaultStaticDir() string {
	root, _ := path.Split(exeFile)
	return path.Join(root, "static")
}

// If this request corresponds to a static file return its path
func (s *Server) findFile(req *http.Request) string {
	//try to serve static files
	staticDirs := s.Config.StaticDirs
	if len(staticDirs) == 0 {
		staticDirs = []string{defaultStaticDir()}
	}
	for _, staticDir := range staticDirs {
		staticFile := path.Join(staticDir, req.URL.Path)
		if fileExists(staticFile) && (req.Method == "GET" || req.Method == "HEAD") {
			return staticFile
		}
	}

	// Try to serve index.html || index.htm
	indexFilenames := []string{"index.html", "index.htm"}
	for _, staticDir := range staticDirs {
		for _, indexFilename := range indexFilenames {
			if indexPath := path.Join(path.Join(staticDir, req.URL.Path), indexFilename); fileExists(indexPath) {
				return indexPath
			}
		}
	}
	return ""
}

// Fully clothed request handler
func (s *Server) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	ctx := &Context{
		Request:  req,
		RawBody:  nil,
		Params:   map[string]string{},
		Server:   s,
		Response: &ResponseWriter{ResponseWriter: w, BodyWriter: w},
		User:     s.User,
	}

	//log the request
	if s.Config.ColorOutput {
		s.Logger.Printf("\033[32;1m%s %s\033[0m", req.Method, req.URL.Path)
	} else {
		s.Logger.Printf("%s %s", req.Method, req.URL.Path)
	}

	//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]
		}
		if s.Config.ColorOutput {
			s.Logger.Printf("\033[37;1mParams: %v\033[0m\n", ctx.Params)
		} else {
			s.Logger.Printf("Params: %v\n", ctx.Params)
		}

	}

	var simpleh SimpleHandler
	route, match := findMatchingRoute(req, s.routes)
	if route != nil {
		if route.method == "WEBSOCKET" {
			// Wrap websocket handler
			openh := func(ctx *Context, args ...string) (err error) {
				// yo dawg we heard you like wrapped functions
				websocket.Handler(func(ws *websocket.Conn) {
					ctx.WebsockConn = ws
					err = route.handler(ctx, args...)
				}).ServeHTTP(ctx.Response, req)
				return err
			}
			simpleh = closeHandler(openh, match[1:]...)
		} else {
			simpleh = closeHandler(route.handler, match[1:]...)
		}
	} else if path := s.findFile(req); path != "" {
		// no custom handler found but there is a file with this name
		simpleh = func(ctx *Context) error {
			http.ServeFile(ctx.Response, ctx.Request, path)
			return nil
		}
	} else {
		// hopeless, 404
		simpleh = func(ctx *Context) error {
			return WebError{404, "Page not found"}
		}
	}
	for _, wrap := range s.Wrappers {
		simpleh = wrapHandler(wrap, simpleh)
	}
	s.applyHandler(simpleh, ctx)
	return
}

func webTime(t time.Time) string {
	ftime := t.Format(time.RFC1123)
	if strings.HasSuffix(ftime, "UTC") {
		ftime = ftime[0:len(ftime)-3] + "GMT"
	}
	return ftime
}

func NewServer() *Server {
	conf := ServerConfig{
		RecoverPanic: true,
		Cert:         "",
		Key:          "",
		ColorOutput:  true,
	}
	s := &Server{
		Config: conf,
		Logger: log.New(os.Stdout, "", log.Ldate|log.Ltime),
		Env:    map[string]interface{}{},
	}
	// Set some default headers
	s.AddWrapper(func(h SimpleHandler, ctx *Context) error {
		ctx.Header().Set("Server", "web.go")
		tm := time.Now().UTC()
		ctx.Header().Set("Date", webTime(tm))
		return h(ctx)
	})
	return s
}

// Package wide proxy functions for global web server object

// Stop the global web server
func Close() {
	mainServer.Close()
}

// Set a logger to be used by the global web server
func SetLogger(logger *log.Logger) {
	mainServer.Logger = logger
}

// The global web server as an object implementing the http.Handler interface
func GetHTTPHandler() http.Handler {
	return mainServer
}

func init() {
	// find the location of the executable
	arg0 := path.Clean(os.Args[0])
	wd, _ := os.Getwd()
	if strings.HasPrefix(arg0, "/") {
		exeFile = arg0
	} else {
		// TODO For robustness, search each directory in $PATH
		exeFile = path.Join(wd, arg0)
	}
	registerDefaultMimetypes()
}
<|endoftext|>