<commit_before>// Package server implements the Tracer server.
package server

import (
	"strings"
	"time"

	"github.com/tracer/tracer"
)

// A StorageTransportEngine returns an instance of a storage transport.
type StorageTransportEngine func(srv *Server, conf map[string]interface{}) (StorageTransport, error)

// A QueryTransportEngine returns an instance of a query transport.
type QueryTransportEngine func(srv *Server, conf map[string]interface{}) (QueryTransport, error)

// A StorageEngine returns an instance of a storage.
type StorageEngine func(conf map[string]interface{}) (Storage, error)

var storageTransportEngines = map[string]StorageTransportEngine{}
var queryTransportEngines = map[string]QueryTransportEngine{}
var storageEngines = map[string]StorageEngine{}

// RegisterStorageTransport registers a storage transport.
func RegisterStorageTransport(name string, engine StorageTransportEngine) {
	storageTransportEngines[name] = engine
}

// GetStorageTransport returns a storage transport by name.
func GetStorageTransport(name string) (StorageTransportEngine, bool) {
	transport, ok := storageTransportEngines[name]
	return transport, ok
}

// RegisterQueryTransport registers a query transport.
func RegisterQueryTransport(name string, engine QueryTransportEngine) {
	queryTransportEngines[name] = engine
}

// GetQueryTransport returns a query transport by name.
func GetQueryTransport(name string) (QueryTransportEngine, bool) {
	transport, ok := queryTransportEngines[name]
	return transport, ok
}

// RegisterStorage registers a storage engine.
func RegisterStorage(name string, engine StorageEngine) {
	storageEngines[name] = engine
}

// GetStorage returns a storage engine by name.
func GetStorage(name string) (StorageEngine, bool) {
	storer, ok := storageEngines[name]
	return storer, ok
}

// A StorageTransport accepts spans via some protocol and sends them
// to a Storer.
type StorageTransport interface {
	// Start starts the transport.
	Start() error
}

// QueryTransport accepts requests via some protocol and answers them.
type QueryTransport interface {
	// Start starts the transport.
	Start() error
}

// Storage allows storing and querying spans.
type Storage interface {
	tracer.Storer
	Queryer
}

// A Purger can delete all traces starting before a certain date.
type Purger interface {
	Purge(before time.Time) error
}

// A Queryer is a backend that allows fetching traces and spans by ID
// or via a more advanced query.
type Queryer interface {
	// TraceByID returns a trace with a specific ID.
	TraceByID(id uint64) (tracer.RawTrace, error)
	// SpanByID returns a span with a specific ID.
	SpanByID(id uint64) (tracer.RawSpan, error)
	// QueryTraces returns all traces that match a query.
	QueryTraces(q Query) ([]tracer.RawTrace, error)

	Services() ([]string, error)
	// TODO(dh): The current Spans function only really exists because
	// of the zipkin UI. We might want a QuerySpans instead.
	Spans(service string) ([]string, error)
	Dependencies() ([]Dependency, error)
}

// Dependency describes the dependency of one service on another.
type Dependency struct {
	Parent string
	Child  string
	Count  uint64
}

// QueryTag describes a single tag or log entry that should be queried
// for.
type QueryTag struct {
	// The key of the tag.
	Key string
	// The value of the tag.
	Value string
	// Whether the value should be checked for.
	CheckValue bool
}

// A Query describes the various conditionals of a query for a trace.
//
// All conditions are ANDed together. Zero values are understood as
// the lack of a constraint.
type Query struct {
	// Only return traces that started at or after this time.
	StartTime time.Time
	// Only return traces that finished before or at this time.
	FinishTime time.Time
	// Only return traces where a span has this operation name.
	OperationName string
	// Only return traces that lasted at least this long.
	MinDuration time.Duration
	// Only return traces that lasted at most this long.
	MaxDuration time.Duration
	// Only return traces where all spans combined have all of these
	// tags.
	AndTags []QueryTag
	// Only return traces where all spans combined have at least one
	// of these tags.
	OrTags []QueryTag
	// How many traces to return. Zero means no limit.
	Num int
}

// Server is an instance of the Tracer application.
type Server struct {
	Storage          Storage
	StorageTransport StorageTransport
	QueryTransports  []QueryTransport
}

type errors struct {
	errs []error
}

func (errs errors) Error() string {
	var s []string
	for _, err := range errs.errs {
		s = append(s, err.Error())
	}
	return strings.Join(s, "\n")
}

func (srv *Server) Start() error {
	errs := make(chan error)
	go func() {
		errs <- srv.StorageTransport.Start()
	}()
	for _, t := range srv.QueryTransports {
		t := t
		go func() {
			errs <- t.Start()
		}()
	}
	var out errors
	for i := 0; i < len(srv.QueryTransports)+1; i++ {
		if err := <-errs; err != nil {
			out.errs = append(out.errs, err)
		}
	}
	if len(out.errs) == 0 {
		return nil
	}
	return out
}
<commit_msg>Document Queryer.Services and Queryer.Dependencies<commit_after>// Package server implements the Tracer server.
package server

import (
	"strings"
	"time"

	"github.com/tracer/tracer"
)

// A StorageTransportEngine returns an instance of a storage transport.
type StorageTransportEngine func(srv *Server, conf map[string]interface{}) (StorageTransport, error)

// A QueryTransportEngine returns an instance of a query transport.
type QueryTransportEngine func(srv *Server, conf map[string]interface{}) (QueryTransport, error)

// A StorageEngine returns an instance of a storage.
type StorageEngine func(conf map[string]interface{}) (Storage, error)

var storageTransportEngines = map[string]StorageTransportEngine{}
var queryTransportEngines = map[string]QueryTransportEngine{}
var storageEngines = map[string]StorageEngine{}

// RegisterStorageTransport registers a storage transport.
func RegisterStorageTransport(name string, engine StorageTransportEngine) {
	storageTransportEngines[name] = engine
}

// GetStorageTransport returns a storage transport by name.
func GetStorageTransport(name string) (StorageTransportEngine, bool) {
	transport, ok := storageTransportEngines[name]
	return transport, ok
}

// RegisterQueryTransport registers a query transport.
func RegisterQueryTransport(name string, engine QueryTransportEngine) {
	queryTransportEngines[name] = engine
}

// GetQueryTransport returns a query transport by name.
func GetQueryTransport(name string) (QueryTransportEngine, bool) {
	transport, ok := queryTransportEngines[name]
	return transport, ok
}

// RegisterStorage registers a storage engine.
func RegisterStorage(name string, engine StorageEngine) {
	storageEngines[name] = engine
}

// GetStorage returns a storage engine by name.
func GetStorage(name string) (StorageEngine, bool) {
	storer, ok := storageEngines[name]
	return storer, ok
}

// A StorageTransport accepts spans via some protocol and sends them
// to a Storer.
type StorageTransport interface {
	// Start starts the transport.
	Start() error
}

// QueryTransport accepts requests via some protocol and answers them.
type QueryTransport interface {
	// Start starts the transport.
	Start() error
}

// Storage allows storing and querying spans.
type Storage interface {
	tracer.Storer
	Queryer
}

// A Purger can delete all traces starting before a certain date.
type Purger interface {
	Purge(before time.Time) error
}

// A Queryer is a backend that allows fetching traces and spans by ID
// or via a more advanced query.
type Queryer interface {
	// TraceByID returns a trace with a specific ID.
	TraceByID(id uint64) (tracer.RawTrace, error)
	// SpanByID returns a span with a specific ID.
	SpanByID(id uint64) (tracer.RawSpan, error)
	// QueryTraces returns all traces that match a query.
	QueryTraces(q Query) ([]tracer.RawTrace, error)

	// Services returns a list of all services.
	Services() ([]string, error)
	// TODO(dh): The current Spans function only really exists because
	// of the zipkin UI. We might want a QuerySpans instead.
	Spans(service string) ([]string, error)
	// Dependencies returns the dependencies between services.
	Dependencies() ([]Dependency, error)
}

// Dependency describes the dependency of one service on another.
type Dependency struct {
	Parent string
	Child  string
	Count  uint64
}

// QueryTag describes a single tag or log entry that should be queried
// for.
type QueryTag struct {
	// The key of the tag.
	Key string
	// The value of the tag.
	Value string
	// Whether the value should be checked for.
	CheckValue bool
}

// A Query describes the various conditionals of a query for a trace.
//
// All conditions are ANDed together. Zero values are understood as
// the lack of a constraint.
type Query struct {
	// Only return traces that started at or after this time.
	StartTime time.Time
	// Only return traces that finished before or at this time.
	FinishTime time.Time
	// Only return traces where a span has this operation name.
	OperationName string
	// Only return traces that lasted at least this long.
	MinDuration time.Duration
	// Only return traces that lasted at most this long.
	MaxDuration time.Duration
	// Only return traces where all spans combined have all of these
	// tags.
	AndTags []QueryTag
	// Only return traces where all spans combined have at least one
	// of these tags.
	OrTags []QueryTag
	// How many traces to return. Zero means no limit.
	Num int
}

// Server is an instance of the Tracer application.
type Server struct {
	Storage          Storage
	StorageTransport StorageTransport
	QueryTransports  []QueryTransport
}

type errors struct {
	errs []error
}

func (errs errors) Error() string {
	var s []string
	for _, err := range errs.errs {
		s = append(s, err.Error())
	}
	return strings.Join(s, "\n")
}

func (srv *Server) Start() error {
	errs := make(chan error)
	go func() {
		errs <- srv.StorageTransport.Start()
	}()
	for _, t := range srv.QueryTransports {
		t := t
		go func() {
			errs <- t.Start()
		}()
	}
	var out errors
	for i := 0; i < len(srv.QueryTransports)+1; i++ {
		if err := <-errs; err != nil {
			out.errs = append(out.errs, err)
		}
	}
	if len(out.errs) == 0 {
		return nil
	}
	return out
}
<|endoftext|>


<commit_before>// Server is the base tcpez server. It sets up a tcp listener on an address,
// and given a RequestHandler, it parses the tcpez protocol format and turns
// it into individual request/responses. Each connection is handled on a
// seperate goroutine and pipelined requests are first parsed then farmed
// to seperate goroutines. Pipelined requests from the client are handled
// seamlessly this way, each seperate request is passed to its own RequestHandler
// with its own Span.
package tcpez

import (
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/op/go-logging"
	"io"
	"net"
	"sync"
)

var log = logging.MustGetLogger("tcpez")
var LogFormat = logging.MustStringFormatter("%{time:2006-01-02T15:04:05.999999999Z07:00} %{level} [%{module}] %{message}")

func init() {
	logging.SetLevel(logging.INFO, "tcpez")
	logging.SetFormatter(LogFormat)
}

// Server is the base struct that wraps the tcp listener and allows
// for setting the RequestHandler that takes each request and returns
// an response
type Server struct {
	// Address is the string address that the tcp listener is
	// bound to
	Address string
	// Handler is a value that responds to the RequestHandler
	// interface.
	Handler RequestHandler
	// StatsRecorder is the value that delivers stats to a collection
	// agent. This is the DebugStatsRecorder by default (does nothing)
	// but can be swapped to send to StatsD. This is passed to each
	// Span created and passed to the RequestHandler
	Stats StatsRecorder
	// UUIDGenerator generates UUIDs for each Span, by default this uses
	// a simple hash function but can be swapped out for something more
	// complex (a vector-clock style UUID generator for example)
	UUIDGenerator UUIDGenerator

	conn     *net.TCPListener
	isClosed bool
}

// RequestHandler is the basic interface for setting up the request handling
// logic of a tcpez server. The server handles all the request parsing and setup
// as well as the response encoding. All you have to do to create a working server
// is create an object that has a Respond() method that takes a byte slice (which
// is the request) and a Span (which allows you to track timings and meta data
// through the request) and then it returns a byte slice which is the response.
//
//        type MyHandler struct
//
//        func (h *MyHandler) Respond(req []byte, span *tcpez.Span) (res []byte, err error) {
//              if string(req) == "PING" {
//                  span.Attr("command", "PING")
//                  return []byte{"PONG"}, nil
//              }
//        }
//
type RequestHandler interface {
	Respond([]byte, *Span) ([]byte, error)
}

// NewServer is the tcpez server intializer. It only requires two parameters,
// an address to bind to (same format as net.ListenTCP) and a RequestHandler
// which serves the requests.
func NewServer(address string, handler RequestHandler) (s *Server, err error) {
	tcpAddr, err := net.ResolveTCPAddr("tcp", address)
	if err != nil {
		return nil, err
	}
	l, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return nil, err
	}

	return &Server{Address: address, conn: l, Handler: handler, Stats: new(DebugStatsRecorder), UUIDGenerator: DefaultUUIDGenerator}, nil
}

// Start starts the connection handling and request processing loop.
// This is a blocking operation and can be started in a goroutine.
func (s *Server) Start() {
	log.Debug("Listening on %s", s.conn.Addr().String())
	for {
		if s.isClosed == true {
			break
		}
		clientConn, err := s.conn.Accept()
		if err != nil {
			log.Warning(err.Error())
			break
		}
		go s.handle(clientConn)
	}
	log.Debug("Closing %s", s.conn.Addr().String())
}

// Close closes the server listener to any more connections
func (s *Server) Close() (err error) {
	if s.isClosed == false {
		err = s.conn.Close()
		s.isClosed = true
		return
	}
	return errors.New("Closing already closed connection")
}

func (s *Server) handle(clientConn net.Conn) {
	log.Debug("[tcpez] New client(%s)", clientConn.RemoteAddr())
	for {
		header, response, err := s.readHeaderAndHandleRequest(clientConn)
		if err != nil {
			if err == io.EOF {
				// EOF the client has disconnected
				break
			}
			log.Error(err.Error())
			s.Stats.Increment("operation.failure")
			return
		}
		err = s.sendResponse(clientConn, header, response)
		if err != nil {
			log.Error(err.Error())
			s.Stats.Increment("operation.failure")
			return
		}
		s.Stats.Increment("operation.success")
	}
}

func (s *Server) readHeaderAndHandleRequest(buf io.Reader) (header int32, response []byte, err error) {
	var size int32
	err = binary.Read(buf, binary.BigEndian, &size)
	if err != nil {
		return 0, nil, err
	}
	if size < 0 {
		// this is a pipelined request
		var wg sync.WaitGroup
		count := -size
		requests := make([][]byte, count)
		responses := make([][]byte, count)
		for r := 0; int32(r) < count; r++ {
			request, err := s.parseRequest(buf, 0)
			if err == nil {
				requests[r] = request
				wg.Add(1)
				go func(index int) {
					res, err := s.handleRequest(requests[index], true)
					if err == nil {
						responses[index] = res
					}
					wg.Done()
				}(r)
			}
		}
		wg.Wait()
		output := bytes.NewBuffer(nil)
		for j := 0; int32(j) < count; j++ {
			length := int32(len(responses[j]))
			err = binary.Write(output, binary.BigEndian, length)
			if err == nil {
				output.Write(responses[j])
			}
		}
		return int32(-count), output.Bytes(), err
	} else {
		request, err := s.parseRequest(buf, size)
		if err != nil {
			return 0, nil, err
		}
		response, err := s.handleRequest(request, false)
		if err != nil {
			return 0, nil, err
		}
		return int32(len(response)), response, nil
	}
	return
}

func (s *Server) sendResponse(w io.Writer, header int32, data []byte) (err error) {
	err = binary.Write(w, binary.BigEndian, header)
	if err != nil {
		return err
	}
	_, err = w.Write(data)
	return
}

func (s *Server) parseRequest(buf io.Reader, size int32) (request []byte, err error) {
	if size == int32(0) {
		err = binary.Read(buf, binary.BigEndian, &size)
		if err != nil {
			return nil, err
		}
	}
	request = make([]byte, size)
	_, err = io.ReadFull(buf, request)
	log.Debug("Server Request: %s", request)
	if err != nil {
		return nil, err
	}
	return
}

func (s *Server) handleRequest(request []byte, multi bool) (response []byte, err error) {
	span := NewSpan(s.UUIDGenerator())
	if multi == true {
		span.Attr("multi", "true")
	}
	span.Stats = s.Stats
	span.Start("duration")
	response, err = s.Handler.Respond(request, span)
	span.Finish("duration")
	log.Info("%s", span.JSON())
	span.Record()
	return
}
<commit_msg>Implement read timeouts on client connections (at 5 mins)<commit_after>// Server is the base tcpez server. It sets up a tcp listener on an address,
// and given a RequestHandler, it parses the tcpez protocol format and turns
// it into individual request/responses. Each Connection is handled on a
// seperate goroutine and pipelined requests are first parsed then farmed
// to seperate goroutines. Pipelined requests from the client are handled
// seamlessly this way, each seperate request is passed to its own RequestHandler
// with its own Span.
package tcpez

import (
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/op/go-logging"
	"io"
	"net"
	"sync"
	"time"
)

var log = logging.MustGetLogger("tcpez")
var LogFormat = logging.MustStringFormatter("%{time:2006-01-02T15:04:05.999999999Z07:00} %{level} [%{module}] %{message}")

func init() {
	logging.SetLevel(logging.INFO, "tcpez")
	logging.SetFormatter(LogFormat)
}

// Server is the base struct that wraps the tcp listener and allows
// for setting the RequestHandler that takes each request and returns
// an response
type Server struct {
	// Address is the string address that the tcp listener is
	// bound to
	Address string
	// Handler is a value that responds to the RequestHandler
	// interface.
	Handler RequestHandler
	// StatsRecorder is the value that delivers stats to a collection
	// agent. This is the DebugStatsRecorder by default (does nothing)
	// but can be swapped to send to StatsD. This is passed to each
	// Span created and passed to the RequestHandler
	Stats StatsRecorder
	// UUIDGenerator generates UUIDs for each Span, by default this uses
	// a simple hash function but can be swapped out for something more
	// complex (a vector-clock style UUID generator for example)
	UUIDGenerator UUIDGenerator

	// The underlying TCP listener, access if you need to set timeouts, etc
	Conn     *net.TCPListener
	isClosed bool
}

// RequestHandler is the basic interface for setting up the request handling
// logic of a tcpez server. The server handles all the request parsing and setup
// as well as the response encoding. All you have to do to create a working server
// is create an object that has a Respond() method that takes a byte slice (which
// is the request) and a Span (which allows you to track timings and meta data
// through the request) and then it returns a byte slice which is the response.
//
//        type MyHandler struct
//
//        func (h *MyHandler) Respond(req []byte, span *tcpez.Span) (res []byte, err error) {
//              if string(req) == "PING" {
//                  span.Attr("command", "PING")
//                  return []byte{"PONG"}, nil
//              }
//        }
//
type RequestHandler interface {
	Respond([]byte, *Span) ([]byte, error)
}

// NewServer is the tcpez server intializer. It only requires two parameters,
// an address to bind to (same format as net.ListenTCP) and a RequestHandler
// which serves the requests.
func NewServer(address string, handler RequestHandler) (s *Server, err error) {
	tcpAddr, err := net.ResolveTCPAddr("tcp", address)
	if err != nil {
		return nil, err
	}
	l, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return nil, err
	}

	return &Server{Address: address, Conn: l, Handler: handler, Stats: new(DebugStatsRecorder), UUIDGenerator: DefaultUUIDGenerator}, nil
}

// Start starts the Connection handling and request processing loop.
// This is a blocking operation and can be started in a goroutine.
func (s *Server) Start() {
	log.Debug("Listening on %s", s.Conn.Addr().String())
	for {
		if s.isClosed == true {
			break
		}
		clientConn, err := s.Conn.Accept()
		if err != nil {
			log.Warning(err.Error())
			break
		}
		go s.handle(clientConn)
	}
	log.Debug("Closing %s", s.Conn.Addr().String())
}

// Close closes the server listener to any more Connections
func (s *Server) Close() (err error) {
	if s.isClosed == false {
		err = s.Conn.Close()
		s.isClosed = true
		return
	}
	return errors.New("Closing already closed Connection")
}

func (s *Server) handle(clientConn net.Conn) {
	log.Debug("[tcpez] New client(%s)", clientConn.RemoteAddr())
	for {
		// Timeout the connection after 5 mins
		clientConn.SetReadDeadline(time.Now().Add(5 * time.Minute))
		header, response, err := s.readHeaderAndHandleRequest(clientConn)
		if err != nil {
			if err == io.EOF || err.(net.Error).Timeout() == true {
				// EOF the client has disConnected
				break
			}
			log.Error(err.Error())
			s.Stats.Increment("operation.failure")
			return
		}
		err = s.sendResponse(clientConn, header, response)
		if err != nil {
			log.Error(err.Error())
			s.Stats.Increment("operation.failure")
			return
		}
		s.Stats.Increment("operation.success")
	}
}

func (s *Server) readHeaderAndHandleRequest(buf io.Reader) (header int32, response []byte, err error) {
	var size int32
	err = binary.Read(buf, binary.BigEndian, &size)
	if err != nil {
		return 0, nil, err
	}
	if size < 0 {
		// this is a pipelined request
		var wg sync.WaitGroup
		count := -size
		requests := make([][]byte, count)
		responses := make([][]byte, count)
		for r := 0; int32(r) < count; r++ {
			request, err := s.parseRequest(buf, 0)
			if err == nil {
				requests[r] = request
				wg.Add(1)
				go func(index int) {
					res, err := s.handleRequest(requests[index], true)
					if err == nil {
						responses[index] = res
					}
					wg.Done()
				}(r)
			}
		}
		wg.Wait()
		output := bytes.NewBuffer(nil)
		for j := 0; int32(j) < count; j++ {
			length := int32(len(responses[j]))
			err = binary.Write(output, binary.BigEndian, length)
			if err == nil {
				output.Write(responses[j])
			}
		}
		return int32(-count), output.Bytes(), err
	} else {
		request, err := s.parseRequest(buf, size)
		if err != nil {
			return 0, nil, err
		}
		response, err := s.handleRequest(request, false)
		if err != nil {
			return 0, nil, err
		}
		return int32(len(response)), response, nil
	}
	return
}

func (s *Server) sendResponse(w io.Writer, header int32, data []byte) (err error) {
	err = binary.Write(w, binary.BigEndian, header)
	if err != nil {
		return err
	}
	_, err = w.Write(data)
	return
}

func (s *Server) parseRequest(buf io.Reader, size int32) (request []byte, err error) {
	if size == int32(0) {
		err = binary.Read(buf, binary.BigEndian, &size)
		if err != nil {
			return nil, err
		}
	}
	request = make([]byte, size)
	_, err = io.ReadFull(buf, request)
	log.Debug("Server Request: %s", request)
	if err != nil {
		return nil, err
	}
	return
}

func (s *Server) handleRequest(request []byte, multi bool) (response []byte, err error) {
	span := NewSpan(s.UUIDGenerator())
	if multi == true {
		span.Attr("multi", "true")
	}
	span.Stats = s.Stats
	span.Start("duration")
	response, err = s.Handler.Respond(request, span)
	span.Finish("duration")
	log.Info("%s", span.JSON())
	span.Record()
	return
}
<|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 (
	"errors"
	"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() error {
	if s.l != nil {
		return s.l.Close()
	}
	return errors.New("closing non-listening web.go server")
}

// 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() error {
	return mainServer.Close()
}

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

func AddWrapper(wrap Wrapper) {
	mainServer.AddWrapper(wrap)
}

// 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>Revert "Do not set content-type HTML 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 (
	"errors"
	"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() error {
	if s.l != nil {
		return s.l.Close()
	}
	return errors.New("closing non-listening web.go server")
}

// 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 {
			// Set the default content-type
			ctx.ContentType("text/html; charset=utf-8")
			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() error {
	return mainServer.Close()
}

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

func AddWrapper(wrap Wrapper) {
	mainServer.AddWrapper(wrap)
}

// 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|>


<commit_before>// Copyright 2016 Jacob Taylor jacob.taylor@gmail.com
// License: Apache2
package main

import (
    "fmt"
    "net"
    "../utils"
    "bufio"
    "encoding/binary"
    "time"
    "os"
    "bytes"
    "io"
    "io/ioutil"
    "log"
)

const nbd_folder = "/sample_disks/"

var characters_per_line = 100
var newline = 0
var line_number = 0

func send_export_list_item(output *bufio.Writer, export_name string) {
    data := make([]byte, 1024)
    length := len(export_name)
    offset := 0

    // length of export name
    binary.BigEndian.PutUint32(data[offset:], uint32(length))  // length of string
    offset += 4

    // export name
    copy(data[offset:], export_name)
    offset += length

    reply_type := uint32(2)     // reply_type: NBD_REP_SERVER
    send_message(output, reply_type, uint32(offset), data)
}

func send_ack(output *bufio.Writer) {
    send_message(output, utils.NBD_COMMAND_ACK, 0, nil)
}

func export_name(output *bufio.Writer, conn net.Conn, payload_size int, payload []byte, pad_with_zeros bool) {
    fmt.Printf("have request to bind to: %s\n", string(payload[:payload_size]))

    defer conn.Close()

    var filename bytes.Buffer
    current_directory, err := os.Getwd()
    utils.ErrorCheck(err)
    filename.WriteString(current_directory)
    filename.WriteString(nbd_folder)
    filename.Write(payload[:payload_size])

    fmt.Printf("Opening file: %s\n", filename.String())

    file, err := os.OpenFile(filename.String(), os.O_RDWR, 0644)

    utils.ErrorCheck(err)
    if err != nil {
        return
    }

    buffer := make([]byte, 256)
    offset := 0

    fs, err := file.Stat()
    file_size := uint64(fs.Size())

    binary.BigEndian.PutUint64(buffer[offset:], file_size)  // size
    offset += 8

    binary.BigEndian.PutUint16(buffer[offset:], 1)  // flags
    offset += 2

    if pad_with_zeros {
        offset += 124               // pad with 124 zeroes
    }

    _, err = output.Write(buffer[:offset])
    output.Flush()
    utils.ErrorCheck(err)

    buffer = make([]byte, 2048*1024)    // set the buffer to 2mb
    conn_reader := bufio.NewReader(conn)
    for {
        waiting_for := 28       // wait for at least the minimum payload size

        _, err := io.ReadFull(conn_reader, buffer[:waiting_for])
        if err == io.EOF {
            fmt.Printf("Abort detected, escaping processing loop\n")
            break
        }
        utils.ErrorCheck(err)

        //magic := binary.BigEndian.Uint32(buffer)
        command := binary.BigEndian.Uint32(buffer[4:8])
        //handle := binary.BigEndian.Uint64(buffer[8:16])
        from := binary.BigEndian.Uint64(buffer[16:24])
        length := binary.BigEndian.Uint32(buffer[24:28])

        newline += 1;
        if newline % characters_per_line == 0 {
            line_number++
            fmt.Printf("\n%5d: ", line_number * 100)
            newline -= characters_per_line
        }

        switch command {
        case utils.NBD_COMMAND_READ:
            fmt.Printf(".")

            _, err = file.ReadAt(buffer[16:16+length], int64(from))
            utils.ErrorCheck(err)

            binary.BigEndian.PutUint32(buffer[:4], utils.NBD_REPLY_MAGIC)
            binary.BigEndian.PutUint32(buffer[4:8], 0)                      // error bits

            conn.Write(buffer[:16+length])

            continue
        case utils.NBD_COMMAND_WRITE:
            fmt.Printf("W")

            //waiting_for += int(length)                   // wait for the additional payload
            //fmt.Printf("About to read the data that we should be writing out.")
            _, err := io.ReadFull(conn_reader, buffer[28:28+length])
            if err == io.EOF {
                fmt.Printf("Abort detected, escaping processing loop\n")
                break
            }
            utils.ErrorCheck(err)
            //fmt.Printf("Done reading the data that should be written")

            _, err = file.WriteAt(buffer[28:28+length], int64(from))
            utils.ErrorCheck(err)

            file.Sync()

            // let them know we are done
            binary.BigEndian.PutUint32(buffer[:4], utils.NBD_REPLY_MAGIC)
            binary.BigEndian.PutUint32(buffer[4:8], 0)                      // error bits

            conn.Write(buffer[:16])

            continue

        case utils.NBD_COMMAND_DISCONNECT:
            fmt.Printf("D")

            file.Sync()
            return
        }
    }
}

func send_export_list(output *bufio.Writer) {
    current_directory, err := os.Getwd()
    files, err := ioutil.ReadDir(current_directory + nbd_folder)
    if err != nil {
        log.Fatal(err)
    }
    for _, file := range files {
        send_export_list_item(output, file.Name())
    }

    send_ack(output)
}

func send_message(output *bufio.Writer, reply_type uint32, length uint32, data []byte ) {
    endian := binary.BigEndian
    buffer := make([]byte, 1024)
    offset := 0

    endian.PutUint64(buffer[offset:], utils.NBD_SERVER_SEND_REPLY_MAGIC)
    offset += 8

    endian.PutUint32(buffer[offset:], uint32(3))  // not sure what this is....
    offset += 4

    endian.PutUint32(buffer[offset:], reply_type)  // reply_type: NBD_REP_SERVER
    offset += 4

    endian.PutUint32(buffer[offset:], length)  // length of package
    offset += 4

    if data != nil {
        copy(buffer[offset:], data[0:length])
        offset += int(length)
    }

    data_to_send := buffer[:offset]
    output.Write(data_to_send)
    output.Flush()

    utils.LogData("Just sent:", offset, data_to_send)
}

func main() {

    if len(os.Args) <  3 {
        panic("missing arguments:  (ipaddress) (portnumber)")
        return
    }

    listener, err := net.Listen("tcp", os.Args[1] + ":" + os.Args[2])
    utils.ErrorCheck(err)

    fmt.Printf("Hello World, we have %v\n", listener)
    reply_magic := make([]byte, 4)
    binary.BigEndian.PutUint32(reply_magic, utils.NBD_REPLY_MAGIC)

    for {
        conn, err := listener.Accept()
        utils.ErrorCheck(err)

        fmt.Printf("We have a new connection from: %s\n", conn.RemoteAddr())
        output := bufio.NewWriter(conn)

        output.WriteString("NBDMAGIC")      // init password
        output.WriteString("IHAVEOPT")      // Magic
        //fmt.printf("arg ")
        //output.Write([]byte{0, byte(os.Args[3][1])})
        //output.Write([]byte{0, 3})          // Ubuntu
        output.Write([]byte{0, 0})        // Qemu

        output.Flush()

        // Fetch the data until we get the initial options
        data := make([]byte, 1024)
        offset := 0
        waiting_for := 16       // wait for at least the minimum payload size

        packet_count := 0
        for offset < waiting_for {
            length, err := conn.Read(data[offset:])
            if length > 0 {
                packet_count += 1
            }
            offset += length
            utils.ErrorCheck(err)
            //utils.LogData("Reading instruction", offset, data)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
            // If we are requesting an export, make sure we have the length of the data for the export name.
            if offset > 15 && binary.BigEndian.Uint32(data[12:]) == utils.NBD_COMMAND_EXPORT_NAME {
                waiting_for = 20
            }
        }

        fmt.Printf("%d packets processed to get %d bytes\n", packet_count, offset)
        utils.LogData("Received from client", offset, data)
        options := binary.BigEndian.Uint32(data[:4])
        command := binary.BigEndian.Uint32(data[12:])
        payload_size := int(binary.BigEndian.Uint32(data[16:]))
        fmt.Printf("Options are: %v\n", options)
        if (options & utils.NBD_FLAG_FIXED_NEW_STYLE) == utils.NBD_FLAG_FIXED_NEW_STYLE {
            fmt.Printf("Fixed New Style option requested\n")
        }
        pad_with_zeros := true
        if (options & utils.NBD_FLAG_NO_ZEROES) == utils.NBD_FLAG_NO_ZEROES {
            pad_with_zeros = false
            fmt.Printf("No Zero Padding option requested\n")
        }

        fmt.Sprintf("command is: %d\npayload_size is: %d\n", command, payload_size)
        waiting_for += int(payload_size)
        for offset < waiting_for {
            length, err := conn.Read(data[offset:])
            offset += length
            utils.ErrorCheck(err)
            utils.LogData("Reading instruction", offset, data)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
        }
        payload := make([]byte, payload_size)

        if payload_size > 0{
            copy(payload, data[20:])
        }

        utils.LogData("Payload is:", payload_size, payload)
        fmt.Printf("command is: %v\n", command)

        // At this point, we have the command, payload size, and payload.
        switch command {
        case utils.NBD_COMMAND_LIST:
            send_export_list(output)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload, pad_with_zeros)
            break
        }
    }

}
<commit_msg>fixed #13 nicer printing messages on startup and shutdown<commit_after>// Copyright 2016 Jacob Taylor jacob.taylor@gmail.com
// License: Apache2
package main

import (
    "fmt"
    "net"
    "../utils"
    "bufio"
    "encoding/binary"
    "time"
    "os"
    "bytes"
    "io"
    "io/ioutil"
    "log"
)

const nbd_folder = "/sample_disks/"

var characters_per_line = 100
var newline = 0
var line_number = 0

func send_export_list_item(output *bufio.Writer, export_name string) {
    data := make([]byte, 1024)
    length := len(export_name)
    offset := 0

    // length of export name
    binary.BigEndian.PutUint32(data[offset:], uint32(length))  // length of string
    offset += 4

    // export name
    copy(data[offset:], export_name)
    offset += length

    reply_type := uint32(2)     // reply_type: NBD_REP_SERVER
    send_message(output, reply_type, uint32(offset), data)
}

func send_ack(output *bufio.Writer) {
    send_message(output, utils.NBD_COMMAND_ACK, 0, nil)
}

func export_name(output *bufio.Writer, conn net.Conn, payload_size int, payload []byte, pad_with_zeros bool) {
    fmt.Printf("have request to bind to: %s\n", string(payload[:payload_size]))

    defer conn.Close()

    var filename bytes.Buffer
    current_directory, err := os.Getwd()
    utils.ErrorCheck(err)
    filename.WriteString(current_directory)
    filename.WriteString(nbd_folder)
    filename.Write(payload[:payload_size])

    fmt.Printf("Opening file: %s\n", filename.String())

    file, err := os.OpenFile(filename.String(), os.O_RDWR, 0644)

    utils.ErrorCheck(err)
    if err != nil {
        return
    }

    buffer := make([]byte, 256)
    offset := 0

    fs, err := file.Stat()
    file_size := uint64(fs.Size())

    binary.BigEndian.PutUint64(buffer[offset:], file_size)  // size
    offset += 8

    binary.BigEndian.PutUint16(buffer[offset:], 1)  // flags
    offset += 2

    if pad_with_zeros {
        offset += 124               // pad with 124 zeroes
    }

    _, err = output.Write(buffer[:offset])
    output.Flush()
    utils.ErrorCheck(err)

    buffer = make([]byte, 2048*1024)    // set the buffer to 2mb
    conn_reader := bufio.NewReader(conn)
    for {
        waiting_for := 28       // wait for at least the minimum payload size

        _, err := io.ReadFull(conn_reader, buffer[:waiting_for])
        if err == io.EOF {
            fmt.Printf("Abort detected, escaping processing loop\n")
            break
        }
        utils.ErrorCheck(err)

        //magic := binary.BigEndian.Uint32(buffer)
        command := binary.BigEndian.Uint32(buffer[4:8])
        //handle := binary.BigEndian.Uint64(buffer[8:16])
        from := binary.BigEndian.Uint64(buffer[16:24])
        length := binary.BigEndian.Uint32(buffer[24:28])

        newline += 1;
        if newline % characters_per_line == 0 {
            line_number++
            fmt.Printf("\n%5d: ", line_number * 100)
            newline -= characters_per_line
        }

        switch command {
        case utils.NBD_COMMAND_READ:
            fmt.Printf(".")

            _, err = file.ReadAt(buffer[16:16+length], int64(from))
            utils.ErrorCheck(err)

            binary.BigEndian.PutUint32(buffer[:4], utils.NBD_REPLY_MAGIC)
            binary.BigEndian.PutUint32(buffer[4:8], 0)                      // error bits

            conn.Write(buffer[:16+length])

            continue
        case utils.NBD_COMMAND_WRITE:
            fmt.Printf("W")

            //waiting_for += int(length)                   // wait for the additional payload
            //fmt.Printf("About to read the data that we should be writing out.")
            _, err := io.ReadFull(conn_reader, buffer[28:28+length])
            if err == io.EOF {
                fmt.Printf("Abort detected, escaping processing loop\n")
                break
            }
            utils.ErrorCheck(err)
            //fmt.Printf("Done reading the data that should be written")

            _, err = file.WriteAt(buffer[28:28+length], int64(from))
            utils.ErrorCheck(err)

            file.Sync()

            // let them know we are done
            binary.BigEndian.PutUint32(buffer[:4], utils.NBD_REPLY_MAGIC)
            binary.BigEndian.PutUint32(buffer[4:8], 0)                      // error bits

            conn.Write(buffer[:16])

            continue

        case utils.NBD_COMMAND_DISCONNECT:
            fmt.Printf("D")

            file.Sync()
            return
        }
    }
}

func send_export_list(output *bufio.Writer) {
    current_directory, err := os.Getwd()
    files, err := ioutil.ReadDir(current_directory + nbd_folder)
    if err != nil {
        log.Fatal(err)
    }
    for _, file := range files {
        send_export_list_item(output, file.Name())
    }

    send_ack(output)
}

func send_message(output *bufio.Writer, reply_type uint32, length uint32, data []byte ) {
    endian := binary.BigEndian
    buffer := make([]byte, 1024)
    offset := 0

    endian.PutUint64(buffer[offset:], utils.NBD_SERVER_SEND_REPLY_MAGIC)
    offset += 8

    endian.PutUint32(buffer[offset:], uint32(3))  // not sure what this is....
    offset += 4

    endian.PutUint32(buffer[offset:], reply_type)  // reply_type: NBD_REP_SERVER
    offset += 4

    endian.PutUint32(buffer[offset:], length)  // length of package
    offset += 4

    if data != nil {
        copy(buffer[offset:], data[0:length])
        offset += int(length)
    }

    data_to_send := buffer[:offset]
    output.Write(data_to_send)
    output.Flush()

    utils.LogData("Just sent:", offset, data_to_send)
}

func main() {

    if len(os.Args) <  3 {
        panic("missing arguments:  (ipaddress) (portnumber)")
        return
    }

    listener, err := net.Listen("tcp", os.Args[1] + ":" + os.Args[2])
    utils.ErrorCheck(err)

    fmt.Printf("ABlox server online\n")

    reply_magic := make([]byte, 4)
    binary.BigEndian.PutUint32(reply_magic, utils.NBD_REPLY_MAGIC)

    defer fmt.Printf("End of line\n")

    for {
        conn, err := listener.Accept()
        utils.ErrorCheck(err)

        fmt.Printf("We have a new connection from: %s\n", conn.RemoteAddr())
        output := bufio.NewWriter(conn)

        output.WriteString("NBDMAGIC")      // init password
        output.WriteString("IHAVEOPT")      // Magic
        //fmt.printf("arg ")
        //output.Write([]byte{0, byte(os.Args[3][1])})
        //output.Write([]byte{0, 3})          // Ubuntu
        output.Write([]byte{0, 0})        // Qemu

        output.Flush()

        // Fetch the data until we get the initial options
        data := make([]byte, 1024)
        offset := 0
        waiting_for := 16       // wait for at least the minimum payload size

        packet_count := 0
        for offset < waiting_for {
            length, err := conn.Read(data[offset:])
            if length > 0 {
                packet_count += 1
            }
            offset += length
            utils.ErrorCheck(err)
            //utils.LogData("Reading instruction", offset, data)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
            // If we are requesting an export, make sure we have the length of the data for the export name.
            if offset > 15 && binary.BigEndian.Uint32(data[12:]) == utils.NBD_COMMAND_EXPORT_NAME {
                waiting_for = 20
            }
        }

        fmt.Printf("%d packets processed to get %d bytes\n", packet_count, offset)
        utils.LogData("Received from client", offset, data)
        options := binary.BigEndian.Uint32(data[:4])
        command := binary.BigEndian.Uint32(data[12:])
        payload_size := int(binary.BigEndian.Uint32(data[16:]))
        fmt.Printf("Options are: %v\n", options)
        if (options & utils.NBD_FLAG_FIXED_NEW_STYLE) == utils.NBD_FLAG_FIXED_NEW_STYLE {
            fmt.Printf("Fixed New Style option requested\n")
        }
        pad_with_zeros := true
        if (options & utils.NBD_FLAG_NO_ZEROES) == utils.NBD_FLAG_NO_ZEROES {
            pad_with_zeros = false
            fmt.Printf("No Zero Padding option requested\n")
        }

        fmt.Sprintf("command is: %d\npayload_size is: %d\n", command, payload_size)
        waiting_for += int(payload_size)
        for offset < waiting_for {
            length, err := conn.Read(data[offset:])
            offset += length
            utils.ErrorCheck(err)
            utils.LogData("Reading instruction", offset, data)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
        }
        payload := make([]byte, payload_size)

        if payload_size > 0{
            copy(payload, data[20:])
        }

        utils.LogData("Payload is:", payload_size, payload)
        fmt.Printf("command is: %v\n", command)

        // At this point, we have the command, payload size, and payload.
        switch command {
        case utils.NBD_COMMAND_LIST:
            send_export_list(output)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload, pad_with_zeros)
            break
        }
    }

}
<|endoftext|>


<commit_before>package ostent
import (
	"share/assets"

	"os"
	"fmt"
	"log"
	"net"
	"flag"
	"bytes"
	"strings"
	"net/http"

	"github.com/justinas/alice"
)

type bindValue struct {
	string
	defport string // const
	Host, Port string // available after flag.Parse()
}

func newBind(defstring, defport string) bindValue {
	bv := bindValue{defport: defport}
	bv.Set(defstring)
	return bv
}

// satisfying flag.Value interface
func(bv bindValue) String() string { return string(bv.string); }
func(bv *bindValue) Set(input string) error {
	if input == "" {
		return nil
	}
	if !strings.Contains(input, ":") {
		input = ":" + input
	}
	var err error
	bv.Host, bv.Port, err = net.SplitHostPort(input)
	if err != nil {
		return err
	}
	if bv.Host == "*" {
		bv.Host = ""
	} else if bv.Port == "127" {
		bv.Host = "127.0.0.1"
		bv.Port = bv.defport
	}
	if _, err = net.LookupPort("tcp", bv.Port); err != nil {
		if bv.Host != "" {
			return err
		}
		bv.Host, bv.Port = bv.Port, bv.defport
	}

	bv.string = bv.Host + ":" + bv.Port
	return nil
}

var OstentBindFlag   = newBind(":8050", "8050")
// var CollectdBindFlag = newBind("",      "8051") // "" by default meaning DO NOT BIND
func init() {
	flag.Var(&OstentBindFlag,   "b",            "Bind address")
	flag.Var(&OstentBindFlag,   "bind",         "Bind address")
	// flag.Var(&CollectdBindFlag, "collectdb",    "Bind address for collectd receiving")
	// flag.Var(&CollectdBindFlag, "collectdbind", "Bind address for collectd receiving")
}

type Muxmap map[string]http.HandlerFunc

var stdaccess *logger // a global, available after Serve call

func Serve(listen net.Listener, production bool, extramap Muxmap) error {
	logger := log.New(os.Stderr, "[ostent] ", 0)
	access := log.New(os.Stdout, "", 0)

	stdaccess = NewLogged(production, access)
	recovery := Recovery(production)

	chain := alice.New(
		stdaccess.Constructor,
		recovery .Constructor,
	)
	mux := NewMux(chain.Then)

	for _, path := range assets.AssetNames() {
		hf := chain.Then(serveContentFunc(path))
		mux.Handle("GET",  "/"+ path, hf)
		mux.Handle("HEAD", "/"+ path, hf)
	}

//	mux.Handle("GET",  "/ws", chain.ThenFunc(slashws)) // that would include stdlogger
	mux.Handle("GET",  "/ws", recovery.ConstructorFunc(slashws)) // slashws uses stdlogger itself

	mux.Handle("GET",  "/",   chain.ThenFunc(index))
	mux.Handle("HEAD", "/",   chain.ThenFunc(index))

	// mux.Handle("GET", "/panic", chain.ThenFunc(func(http.ResponseWriter, *http.Request) { panic(fmt.Errorf("I'm panicing")) }))

	if extramap != nil {
		for path, handler := range extramap {
			for _, METH := range []string{"HEAD", "GET", "POST"} {
				mux.Handle(METH, path, chain.Then(handler))
			}
		}
	}

	banner(listen, logger)

	server := &http.Server{Addr: listen.Addr().String(), Handler: mux}
	return server.Serve(listen)
}

func serveContentFunc(path string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		text, err := assets.Uncompressedasset(path)
		if err != nil {
			panic(err)
		}
		modtime, err := assets.ModTime("assets", path)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		http.ServeContent(w, r, path, modtime, bytes.NewReader(text))
	}
}

func banner(listen net.Listener, logger *log.Logger) {
	hostname, _ := getHostname()
	logger.Printf("   %s\n", strings.Repeat("-", len(hostname) + 7))
	if len(hostname) > 19 {
		hostname = hostname[:16] +"..."
	}
	logger.Printf(" / %s ostent \\ \n", hostname)
	logger.Printf("+------------------------------+")

	addr := listen.Addr()
	if h, port, err := net.SplitHostPort(addr.String()); err == nil && h == "::" {
		// wildcard bind

		/* _, IP := NewInterfaces()
		logger.Printf("        http://%s", IP) // */
		addrs, err := net.InterfaceAddrs()
		if err == nil {
			fst := true
			for _, a := range addrs {
				ipnet, ok := a.(*net.IPNet)
				if !ok || strings.Contains(ipnet.IP.String(), ":") {
					continue // no IPv6 for now
				}
				f := fmt.Sprintf("http://%s:%s", ipnet.IP.String(), port)
				if len(f) < 28 {
					f += strings.Repeat(" ", 28 - len(f))
				}
				if !fst {
					logger.Printf("|------------------------------|")
				}
				fst = false
				logger.Printf("| %s |", f)
			}
		}
	} else {
		f := fmt.Sprintf("http://%s", addr.String())
		if len(f) < 28 {
			f += strings.Repeat(" ", 28 - len(f))
		}
		logger.Printf("| %s |", f)
	}
	logger.Printf("+------------------------------+")
}

const VERSION = "0.1.8"
<commit_msg>godoc, gofmt<commit_after>package ostent

import (
	"share/assets"

	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"strings"

	"github.com/justinas/alice"
)

type bindValue struct {
	string
	defport string // const
	Host    string // available after flag.Parse()
	Port    string // available after flag.Parse()
}

func newBind(defstring, defport string) bindValue {
	bv := bindValue{defport: defport}
	bv.Set(defstring)
	return bv
}

// satisfying flag.Value interface
func (bv bindValue) String() string { return string(bv.string) }
func (bv *bindValue) Set(input string) error {
	if input == "" {
		return nil
	}
	if !strings.Contains(input, ":") {
		input = ":" + input
	}
	var err error
	bv.Host, bv.Port, err = net.SplitHostPort(input)
	if err != nil {
		return err
	}
	if bv.Host == "*" {
		bv.Host = ""
	} else if bv.Port == "127" {
		bv.Host = "127.0.0.1"
		bv.Port = bv.defport
	}
	if _, err = net.LookupPort("tcp", bv.Port); err != nil {
		if bv.Host != "" {
			return err
		}
		bv.Host, bv.Port = bv.Port, bv.defport
	}

	bv.string = bv.Host + ":" + bv.Port
	return nil
}

// OstentBindFlag is a bindValue hoding the ostent bind address.
var OstentBindFlag = newBind(":8050", "8050")

// CollectdBindFlag is a bindValue hoding the ostent collectd bind address.
// var CollectdBindFlag = newBind("", "8051") // "" by default meaning DO NOT BIND
func init() {
	flag.Var(&OstentBindFlag, "b", "short for bind")
	flag.Var(&OstentBindFlag, "bind", "Bind address")
	// flag.Var(&CollectdBindFlag, "collectdb",    "short for collectdbind")
	// flag.Var(&CollectdBindFlag, "collectdbind", "Bind address for collectd receiving")
}

// Muxmap is a type of a map of pattern to HandlerFunc.
type Muxmap map[string]http.HandlerFunc

var stdaccess *logger // a global, available after Serve call

// Serve does http.Serve with the listener l and constructed *TrieServeMux.
// production is passed to logging and recoverying middleware.
// Non-nil extramap is passed to the mux.
// Returns http.Serve result.
func Serve(l net.Listener, production bool, extramap Muxmap) error {
	logger := log.New(os.Stderr, "[ostent] ", 0)
	access := log.New(os.Stdout, "", 0)

	stdaccess = NewLogged(production, access)
	recovery := Recovery(production)

	chain := alice.New(
		stdaccess.Constructor,
		 recovery.Constructor,
	)
	mux := NewMux(chain.Then)

	for _, path := range assets.AssetNames() {
		hf := chain.Then(serveContentFunc(path))
		mux.Handle("GET", "/"+path, hf)
		mux.Handle("HEAD", "/"+path, hf)
	}

	//	chain.ThenFunc(slashws)) handler would include stdlogger
	//  slashws uses stdlogger itself
	mux.Handle("GET", "/ws", recovery.ConstructorFunc(slashws))

	mux.Handle("GET", "/", chain.ThenFunc(index))
	mux.Handle("HEAD", "/", chain.ThenFunc(index))

	/* panics := func(http.ResponseWriter, *http.Request) {
		panic(fmt.Errorf("I'm panicing"))
	}
	mux.Handle("GET", "/panic", chain.ThenFunc(panics)) // */

	if extramap != nil {
		for path, handler := range extramap {
			for _, METH := range []string{"HEAD", "GET", "POST"} {
				mux.Handle(METH, path, chain.Then(handler))
			}
		}
	}

	banner(l, logger)

	server := &http.Server{Addr: l.Addr().String(), Handler: mux}
	return server.Serve(l)
}

func serveContentFunc(path string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		text, err := assets.Uncompressedasset(path)
		if err != nil {
			panic(err)
		}
		modtime, err := assets.ModTime("assets", path)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		http.ServeContent(w, r, path, modtime, bytes.NewReader(text))
	}
}

func banner(listen net.Listener, logger *log.Logger) {
	hostname, _ := getHostname()
	logger.Printf("   %s\n", strings.Repeat("-", len(hostname)+7))
	if len(hostname) > 19 {
		hostname = hostname[:16] + "..."
	}
	logger.Printf(" / %s ostent \\ \n", hostname)
	logger.Printf("+------------------------------+")

	addr := listen.Addr()
	if h, port, err := net.SplitHostPort(addr.String()); err == nil && h == "::" {
		// wildcard bind

		/* _, IP := NewInterfaces()
		logger.Printf("        http://%s", IP) // */
		addrs, err := net.InterfaceAddrs()
		if err == nil {
			fst := true
			for _, a := range addrs {
				ipnet, ok := a.(*net.IPNet)
				if !ok || strings.Contains(ipnet.IP.String(), ":") {
					continue // no IPv6 for now
				}
				f := fmt.Sprintf("http://%s:%s", ipnet.IP.String(), port)
				if len(f) < 28 {
					f += strings.Repeat(" ", 28-len(f))
				}
				if !fst {
					logger.Printf("|------------------------------|")
				}
				fst = false
				logger.Printf("| %s |", f)
			}
		}
	} else {
		f := fmt.Sprintf("http://%s", addr.String())
		if len(f) < 28 {
			f += strings.Repeat(" ", 28-len(f))
		}
		logger.Printf("| %s |", f)
	}
	logger.Printf("+------------------------------+")
}

// VERSION of the latest known release.
// Unused in non-production mode.
// Compared with in main.production.go.
const VERSION = "0.1.8"
<|endoftext|>


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

package server

import (
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"path/filepath"

	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/xlab-si/emmy/config"
	"github.com/xlab-si/emmy/crypto/groups"
	"github.com/xlab-si/emmy/log"
	pb "github.com/xlab-si/emmy/protobuf"
	"github.com/xlab-si/emmy/types"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

var _ pb.ProtocolServer = (*Server)(nil)

type Server struct {
	grpcServer *grpc.Server
	logger     log.Logger
	*sessionManager
	*registrationManager
}

// NewProtocolServer initializes an instance of the Server struct and returns a pointer.
// It performs some default configuration (tracing of gRPC communication and interceptors)
// and registers RPC protocol server with gRPC server. It requires TLS cert and keyfile
// in order to establish a secure channel with clients.
func NewProtocolServer(certFile, keyFile string, logger log.Logger) (*Server, error) {
	logger.Info("Instantiating new protocol server")

	// Register our generic service
	logger.Info("Registering services")

	// Obtain TLS credentials
	creds, err := credentials.NewServerTLSFromFile(certFile, keyFile)
	if err != nil {
		return nil, err
	}

	logger.Infof("Successfully read certificate [%s] and key [%s]", certFile, keyFile)

	sessionManager, err := newSessionManager(config.LoadSessionKeyMinByteLen())
	if err != nil {
		logger.Warning(err)
	}

	registrationManager, err := NewRegistrationManager(config.LoadRegistrationDBAddress())
	if err != nil {
		logger.Error(err)
		return nil, err
	}

	// Allow as much concurrent streams as possible and register a gRPC stream interceptor
	// for logging and monitoring purposes.
	server := &Server{
		grpcServer: grpc.NewServer(
			grpc.Creds(creds),
			grpc.MaxConcurrentStreams(math.MaxUint32),
			grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
		),
		logger:              logger,
		sessionManager:      sessionManager,
		registrationManager: registrationManager,
	}

	// Disable tracing by default, as is used for debugging purposes.
	// The user will be able to turn it on via Server's EnableTracing function.
	grpc.EnableTracing = false

	// Register our services with the supporting gRPC server
	pb.RegisterProtocolServer(server.grpcServer, server)
	pb.RegisterInfoServer(server.grpcServer, server)

	// Initialize gRPC metrics offered by Prometheus package
	grpc_prometheus.Register(server.grpcServer)

	logger.Notice("gRPC Services registered")
	return server, nil
}

// Start configures and starts the protocol server at the requested port.
func (s *Server) Start(port int) error {
	connStr := fmt.Sprintf(":%d", port)
	listener, err := net.Listen("tcp", connStr)
	if err != nil {
		return fmt.Errorf("Could not connect: %v", err)
	}

	// Register Prometheus metrics handler and serve metrics page on the desired endpoint.
	// Metrics are handled via HTTP in a separate goroutine as gRPC requests,
	// as grpc server's performance over HTTP (grpcServer.ServeHTTP) is much worse.
	http.Handle("/metrics", prometheus.Handler())

	// After this, /metrics will be available, along with /debug/requests, /debug/events in
	// case server's EnableTracing function is called.
	go http.ListenAndServe(":8881", nil)

	// From here on, gRPC server will accept connections
	s.logger.Noticef("Emmy server listening for connections on port %d", port)
	s.grpcServer.Serve(listener)
	return nil
}

// Teardown stops the protocol server by gracefully stopping enclosed gRPC server.
func (s *Server) Teardown() {
	s.logger.Notice("Tearing down gRPC server")
	s.grpcServer.GracefulStop()
}

// EnableTracing instructs the gRPC framework to enable its tracing capability, which
// is mainly used for debugging purposes.
// Although this function does not explicitly affect the Server struct, it is wired to Server
// in order to provide a nicer API when setting up the server.
func (s *Server) EnableTracing() {
	grpc.EnableTracing = true
	s.logger.Notice("Enabled gRPC tracing")
}

func (s *Server) send(msg *pb.Message, stream pb.Protocol_RunServer) error {
	if err := stream.Send(msg); err != nil {
		return fmt.Errorf("Error sending message:", err)
	}
	s.logger.Infof("Successfully sent response of type %T", msg.Content)
	s.logger.Debugf("%+v", msg)

	return nil
}

func (s *Server) receive(stream pb.Protocol_RunServer) (*pb.Message, error) {
	resp, err := stream.Recv()
	if err == io.EOF {
		return nil, err
	} else if err != nil {
		return nil, fmt.Errorf("An error ocurred: %v", err)
	}
	s.logger.Infof("Received request of type %T from the stream", resp.Content)
	s.logger.Debugf("%+v", resp)

	return resp, nil
}

func (s *Server) Run(stream pb.Protocol_RunServer) error {
	s.logger.Info("Starting new RPC")

	req, err := s.receive(stream)
	if err != nil {
		return err
	}

	reqClientId := req.ClientId
	reqSchemaType := req.Schema
	reqSchemaVariant := req.SchemaVariant

	// Check whether the client requested a valid schema
	reqSchemaTypeStr, schemaValid := pb.SchemaType_name[int32(reqSchemaType)]
	if !schemaValid {
		return fmt.Errorf("Client [", reqClientId, "] requested invalid schema: %v", reqSchemaType)
	}

	// Check whether the client requested a valid schema variant
	reqSchemaVariantStr, variantValid := pb.SchemaVariant_name[int32(reqSchemaVariant)]
	if !variantValid {
		return fmt.Errorf("Client [ %v ] requested invalid schema variant: %v", reqClientId, reqSchemaVariant)
	}

	s.logger.Noticef("Client [ %v ] requested schema %v, variant %v", reqClientId, reqSchemaTypeStr, reqSchemaVariantStr)

	// Convert Sigma, ZKP or ZKPOK protocol type to a types type
	protocolType := types.ToProtocolType(reqSchemaVariant)
	// This curve will be used for all schemes
	curve := groups.P256

	switch reqSchemaType {
	case pb.SchemaType_PEDERSEN_EC:
		err = s.PedersenEC(curve, stream)
	case pb.SchemaType_PEDERSEN:
		group := config.LoadGroup("pedersen")
		err = s.Pedersen(group, stream)
	case pb.SchemaType_SCHNORR:
		group := config.LoadGroup("schnorr")
		err = s.Schnorr(req, group, protocolType, stream)
	case pb.SchemaType_SCHNORR_EC:
		err = s.SchnorrEC(req, protocolType, stream, curve)
	case pb.SchemaType_CSPAILLIER:
		keyDir := config.LoadKeyDirFromConfig()
		secKeyPath := filepath.Join(keyDir, "cspaillierseckey.txt")
		err = s.CSPaillier(req, secKeyPath, stream)
	case pb.SchemaType_PSEUDONYMSYS_CA:
		err = s.PseudonymsysCA(req, stream)
	case pb.SchemaType_PSEUDONYMSYS_NYM_GEN:
		err = s.PseudonymsysGenerateNym(req, stream)
	case pb.SchemaType_PSEUDONYMSYS_ISSUE_CREDENTIAL:
		err = s.PseudonymsysIssueCredential(req, stream)
	case pb.SchemaType_PSEUDONYMSYS_TRANSFER_CREDENTIAL:
		err = s.PseudonymsysTransferCredential(req, stream)
	case pb.SchemaType_PSEUDONYMSYS_CA_EC:
		err = s.PseudonymsysCAEC(curve, req, stream)
	case pb.SchemaType_PSEUDONYMSYS_NYM_GEN_EC:
		err = s.PseudonymsysGenerateNymEC(curve, req, stream)
	case pb.SchemaType_PSEUDONYMSYS_ISSUE_CREDENTIAL_EC:
		err = s.PseudonymsysIssueCredentialEC(curve, req, stream)
	case pb.SchemaType_PSEUDONYMSYS_TRANSFER_CREDENTIAL_EC:
		err = s.PseudonymsysTransferCredentialEC(curve, req, stream)
	case pb.SchemaType_QR:
		group := config.LoadGroup("pseudonymsys")
		err = s.QR(req, group, stream)
	case pb.SchemaType_QNR:
		qr := config.LoadQRRSA("qrsmall") // only for testing
		err = s.QNR(req, qr, stream)
	}

	if err != nil {
		s.logger.Error("Closing RPC due to previous errors")
		return fmt.Errorf("FAIL: %v", err)
	}

	s.logger.Notice("RPC finished successfully")
	return nil
}
<commit_msg>Update calls to fmt.Errorf in server.go<commit_after>/*
 * Copyright 2017 XLAB d.o.o.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package server

import (
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"path/filepath"

	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/xlab-si/emmy/config"
	"github.com/xlab-si/emmy/crypto/groups"
	"github.com/xlab-si/emmy/log"
	pb "github.com/xlab-si/emmy/protobuf"
	"github.com/xlab-si/emmy/types"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

var _ pb.ProtocolServer = (*Server)(nil)

type Server struct {
	grpcServer *grpc.Server
	logger     log.Logger
	*sessionManager
	*registrationManager
}

// NewProtocolServer initializes an instance of the Server struct and returns a pointer.
// It performs some default configuration (tracing of gRPC communication and interceptors)
// and registers RPC protocol server with gRPC server. It requires TLS cert and keyfile
// in order to establish a secure channel with clients.
func NewProtocolServer(certFile, keyFile string, logger log.Logger) (*Server, error) {
	logger.Info("Instantiating new protocol server")

	// Register our generic service
	logger.Info("Registering services")

	// Obtain TLS credentials
	creds, err := credentials.NewServerTLSFromFile(certFile, keyFile)
	if err != nil {
		return nil, err
	}

	logger.Infof("Successfully read certificate [%s] and key [%s]", certFile, keyFile)

	sessionManager, err := newSessionManager(config.LoadSessionKeyMinByteLen())
	if err != nil {
		logger.Warning(err)
	}

	registrationManager, err := NewRegistrationManager(config.LoadRegistrationDBAddress())
	if err != nil {
		logger.Error(err)
		return nil, err
	}

	// Allow as much concurrent streams as possible and register a gRPC stream interceptor
	// for logging and monitoring purposes.
	server := &Server{
		grpcServer: grpc.NewServer(
			grpc.Creds(creds),
			grpc.MaxConcurrentStreams(math.MaxUint32),
			grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
		),
		logger:              logger,
		sessionManager:      sessionManager,
		registrationManager: registrationManager,
	}

	// Disable tracing by default, as is used for debugging purposes.
	// The user will be able to turn it on via Server's EnableTracing function.
	grpc.EnableTracing = false

	// Register our services with the supporting gRPC server
	pb.RegisterProtocolServer(server.grpcServer, server)
	pb.RegisterInfoServer(server.grpcServer, server)

	// Initialize gRPC metrics offered by Prometheus package
	grpc_prometheus.Register(server.grpcServer)

	logger.Notice("gRPC Services registered")
	return server, nil
}

// Start configures and starts the protocol server at the requested port.
func (s *Server) Start(port int) error {
	connStr := fmt.Sprintf(":%d", port)
	listener, err := net.Listen("tcp", connStr)
	if err != nil {
		return fmt.Errorf("Could not connect: %v", err)
	}

	// Register Prometheus metrics handler and serve metrics page on the desired endpoint.
	// Metrics are handled via HTTP in a separate goroutine as gRPC requests,
	// as grpc server's performance over HTTP (grpcServer.ServeHTTP) is much worse.
	http.Handle("/metrics", prometheus.Handler())

	// After this, /metrics will be available, along with /debug/requests, /debug/events in
	// case server's EnableTracing function is called.
	go http.ListenAndServe(":8881", nil)

	// From here on, gRPC server will accept connections
	s.logger.Noticef("Emmy server listening for connections on port %d", port)
	s.grpcServer.Serve(listener)
	return nil
}

// Teardown stops the protocol server by gracefully stopping enclosed gRPC server.
func (s *Server) Teardown() {
	s.logger.Notice("Tearing down gRPC server")
	s.grpcServer.GracefulStop()
}

// EnableTracing instructs the gRPC framework to enable its tracing capability, which
// is mainly used for debugging purposes.
// Although this function does not explicitly affect the Server struct, it is wired to Server
// in order to provide a nicer API when setting up the server.
func (s *Server) EnableTracing() {
	grpc.EnableTracing = true
	s.logger.Notice("Enabled gRPC tracing")
}

func (s *Server) send(msg *pb.Message, stream pb.Protocol_RunServer) error {
	if err := stream.Send(msg); err != nil {
		return fmt.Errorf("Error sending message: %v", err)
	}
	s.logger.Infof("Successfully sent response of type %T", msg.Content)
	s.logger.Debugf("%+v", msg)

	return nil
}

func (s *Server) receive(stream pb.Protocol_RunServer) (*pb.Message, error) {
	resp, err := stream.Recv()
	if err == io.EOF {
		return nil, err
	} else if err != nil {
		return nil, fmt.Errorf("An error ocurred: %v", err)
	}
	s.logger.Infof("Received request of type %T from the stream", resp.Content)
	s.logger.Debugf("%+v", resp)

	return resp, nil
}

func (s *Server) Run(stream pb.Protocol_RunServer) error {
	s.logger.Info("Starting new RPC")

	req, err := s.receive(stream)
	if err != nil {
		return err
	}

	reqClientId := req.ClientId
	reqSchemaType := req.Schema
	reqSchemaVariant := req.SchemaVariant

	// Check whether the client requested a valid schema
	reqSchemaTypeStr, schemaValid := pb.SchemaType_name[int32(reqSchemaType)]
	if !schemaValid {
		return fmt.Errorf("Client [ %d ] requested invalid schema: %v", reqClientId, reqSchemaType)
	}

	// Check whether the client requested a valid schema variant
	reqSchemaVariantStr, variantValid := pb.SchemaVariant_name[int32(reqSchemaVariant)]
	if !variantValid {
		return fmt.Errorf("Client [ %d ] requested invalid schema variant: %v", reqClientId, reqSchemaVariant)
	}

	s.logger.Noticef("Client [ %v ] requested schema %v, variant %v", reqClientId, reqSchemaTypeStr, reqSchemaVariantStr)

	// Convert Sigma, ZKP or ZKPOK protocol type to a types type
	protocolType := types.ToProtocolType(reqSchemaVariant)
	// This curve will be used for all schemes
	curve := groups.P256

	switch reqSchemaType {
	case pb.SchemaType_PEDERSEN_EC:
		err = s.PedersenEC(curve, stream)
	case pb.SchemaType_PEDERSEN:
		group := config.LoadGroup("pedersen")
		err = s.Pedersen(group, stream)
	case pb.SchemaType_SCHNORR:
		group := config.LoadGroup("schnorr")
		err = s.Schnorr(req, group, protocolType, stream)
	case pb.SchemaType_SCHNORR_EC:
		err = s.SchnorrEC(req, protocolType, stream, curve)
	case pb.SchemaType_CSPAILLIER:
		keyDir := config.LoadKeyDirFromConfig()
		secKeyPath := filepath.Join(keyDir, "cspaillierseckey.txt")
		err = s.CSPaillier(req, secKeyPath, stream)
	case pb.SchemaType_PSEUDONYMSYS_CA:
		err = s.PseudonymsysCA(req, stream)
	case pb.SchemaType_PSEUDONYMSYS_NYM_GEN:
		err = s.PseudonymsysGenerateNym(req, stream)
	case pb.SchemaType_PSEUDONYMSYS_ISSUE_CREDENTIAL:
		err = s.PseudonymsysIssueCredential(req, stream)
	case pb.SchemaType_PSEUDONYMSYS_TRANSFER_CREDENTIAL:
		err = s.PseudonymsysTransferCredential(req, stream)
	case pb.SchemaType_PSEUDONYMSYS_CA_EC:
		err = s.PseudonymsysCAEC(curve, req, stream)
	case pb.SchemaType_PSEUDONYMSYS_NYM_GEN_EC:
		err = s.PseudonymsysGenerateNymEC(curve, req, stream)
	case pb.SchemaType_PSEUDONYMSYS_ISSUE_CREDENTIAL_EC:
		err = s.PseudonymsysIssueCredentialEC(curve, req, stream)
	case pb.SchemaType_PSEUDONYMSYS_TRANSFER_CREDENTIAL_EC:
		err = s.PseudonymsysTransferCredentialEC(curve, req, stream)
	case pb.SchemaType_QR:
		group := config.LoadGroup("pseudonymsys")
		err = s.QR(req, group, stream)
	case pb.SchemaType_QNR:
		qr := config.LoadQRRSA("qrsmall") // only for testing
		err = s.QNR(req, qr, stream)
	}

	if err != nil {
		s.logger.Error("Closing RPC due to previous errors")
		return fmt.Errorf("FAIL: %v", err)
	}

	s.logger.Notice("RPC finished successfully")
	return nil
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"github.com/gansidui/chatserver/handlers"
	"github.com/gansidui/chatserver/packet"
	"github.com/gansidui/chatserver/report"
	"github.com/gansidui/chatserver/utils/convert"
	"github.com/gansidui/chatserver/utils/funcmap"
	"log"
	"net"
	"sync"
	"time"
)

type Server struct {
	exitCh        chan bool        // 结束信号
	waitGroup     *sync.WaitGroup  // 等待goroutine
	funcMap       *funcmap.FuncMap // 映射消息处理函数(uint32 --> func)
	acceptTimeout time.Duration    // 连接超时时间
	readTimeout   time.Duration    // 读超时时间,其实也就是心跳维持时间
	writeTimeout  time.Duration    // 写超时时间
	reqMemPool    *sync.Pool       // 为每个conn分配一个固定的接收缓存
}

func NewServer() *Server {
	return &Server{
		exitCh:        make(chan bool),
		waitGroup:     &sync.WaitGroup{},
		funcMap:       funcmap.NewFuncMap(),
		acceptTimeout: 30,
		readTimeout:   60,
		writeTimeout:  60,
		reqMemPool: &sync.Pool{
			New: func() interface{} {
				return make([]byte, 1024)
			},
		},
	}
}

func (this *Server) SetAcceptTimeout(acceptTimeout time.Duration) {
	this.acceptTimeout = acceptTimeout
}

func (this *Server) SetReadTimeout(readTimeout time.Duration) {
	this.readTimeout = readTimeout
}

func (this *Server) SetWriteTimeout(writeTimeout time.Duration) {
	this.writeTimeout = writeTimeout
}

func (this *Server) Start(listener *net.TCPListener) {
	log.Printf("Start listen on %v\r\n", listener.Addr())
	this.waitGroup.Add(1)
	defer func() {
		listener.Close()
		this.waitGroup.Done()
	}()

	// 防止恶意连接
	go this.dealSpamConn()

	// report记录，定时发送邮件
	go report.Work()

	for {
		select {
		case <-this.exitCh:
			log.Printf("Stop listen on %v\r\n", listener.Addr())
			return
		default:
		}

		listener.SetDeadline(time.Now().Add(this.acceptTimeout))
		conn, err := listener.AcceptTCP()

		if err != nil {
			if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
				// log.Printf("Accept timeout: %v\r\n", opErr)
				continue
			}
			report.AddCount(report.TryConnect, 1)
			log.Printf("Accept error: %v\r\n", err)
			continue
		}

		report.AddCount(report.SuccessConnect, 1)

		// 连接后等待登陆验证
		handlers.ConnMapLoginStatus.Set(conn, time.Now())
		log.Printf("Accept: %v\r\n", conn.RemoteAddr())

		go this.handleClientConn(conn)
	}
}

// 处理恶意连接，定时检测。
// 若conn的loginstatus为nil,则说明conn已经登陆成功。
// 若conn的loginstatus不为nil,则表示loginstatus为该conn连接服务器时的时间戳(time.Time)
// 判断这个时间戳是否已经超过登陆限制时间，若超过，则断开。
func (this *Server) dealSpamConn() {
	limitTime := 60 * time.Second
	ticker := time.NewTicker(limitTime)
	for _ = range ticker.C {
		items := handlers.ConnMapLoginStatus.Items()
		for conn, loginstatus := range items {
			if loginstatus != nil {
				deadline := loginstatus.(time.Time).Add(limitTime)
				if time.Now().After(deadline) {
					conn.(*net.TCPConn).Close()
					handlers.ConnMapLoginStatus.Delete(conn.(*net.TCPConn))
				}
			}
		}

		// 输出当前的连接数
		fmt.Println("当前连接数: ", handlers.ConnMapUuid.Size())
	}
}

func (this *Server) Stop() {
	// close后，所有的exitCh都返回false
	close(this.exitCh)
	this.waitGroup.Wait()
}

func (this *Server) BindMsgHandler(pacType uint32, fn interface{}) error {
	return this.funcMap.Bind(pacType, fn)
}

func (this *Server) handleClientConn(conn *net.TCPConn) {
	this.waitGroup.Add(1)
	defer this.waitGroup.Done()

	receivePackets := make(chan *packet.Packet, 20) // 接收到的包
	chStop := make(chan bool)                       // 通知停止消息处理
	addr := conn.RemoteAddr().String()

	defer func() {
		defer func() {
			if e := recover(); e != nil {
				log.Printf("Panic: %v\r\n", e)
			}
		}()

		handlers.CloseConn(conn)

		log.Printf("Disconnect: %v\r\n", addr)
		chStop <- true
	}()

	// 处理接收到的包
	go this.handlePackets(conn, receivePackets, chStop)

	// 接收数据
	log.Printf("HandleClient: %v\r\n", addr)

	request := this.reqMemPool.Get().([]byte)
	defer this.reqMemPool.Put(request)

	buf := make([]byte, 0)
	var bufLen uint32 = 0

	for {
		select {
		case <-this.exitCh:
			log.Printf("Stop handleClientConn\r\n")
			return
		default:
		}

		conn.SetReadDeadline(time.Now().Add(this.readTimeout))
		readSize, err := conn.Read(request)
		if err != nil {
			log.Printf("Read failed: %v\r\n", err)
			return
		}

		if readSize > 0 {
			buf = append(buf, request[:readSize]...)
			bufLen += uint32(readSize)

			// 包长(4) + 类型(4) + 包体(len([]byte))
			for {
				if bufLen >= 8 {
					pacLen := convert.BytesToUint32(buf[0:4])
					if bufLen >= pacLen {
						receivePackets <- &packet.Packet{
							Len:  pacLen,
							Type: convert.BytesToUint32(buf[4:8]),
							Data: buf[8:pacLen],
						}
						buf = buf[pacLen:]
						bufLen -= pacLen
					} else {
						break
					}
				} else {
					break
				}
			}
		}

	}
}

func (this *Server) handlePackets(conn *net.TCPConn, receivePackets <-chan *packet.Packet, chStop <-chan bool) {
	defer func() {
		if e := recover(); e != nil {
			log.Printf("Panic: %v\r\n", e)
		}
	}()

	for {
		select {
		case <-chStop:
			log.Printf("Stop handle receivePackets.\r\n")
			return

		// 消息包处理
		case p := <-receivePackets:
			// 防止模拟的客户端未经登陆就发送消息
			// 如果接收的不是登陆包，则检查是否已经在线，若没在线，则无视消息包，等待登陆检测机制处理。
			if p.Type != packet.PK_ClientLogin {
				if !handlers.ConnMapUuid.Check(conn) {
					continue
				}
			}

			if this.funcMap.Exist(p.Type) {
				this.funcMap.Call(p.Type, conn, p)
			} else {
				log.Printf("Unknown packet type\r\n")
			}
		}
	}
}
<commit_msg>add ring buffer<commit_after>package server

import (
	"fmt"
	"github.com/gansidui/chatserver/handlers"
	"github.com/gansidui/chatserver/packet"
	"github.com/gansidui/chatserver/report"
	"github.com/gansidui/chatserver/utils/convert"
	"github.com/gansidui/chatserver/utils/funcmap"
	"github.com/gansidui/code/ringbuffer"
	"log"
	"net"
	"sync"
	"time"
)

type Server struct {
	exitCh        chan bool        // 结束信号
	waitGroup     *sync.WaitGroup  // 等待goroutine
	funcMap       *funcmap.FuncMap // 映射消息处理函数(uint32 --> func)
	acceptTimeout time.Duration    // 连接超时时间
	readTimeout   time.Duration    // 读超时时间,其实也就是心跳维持时间
	writeTimeout  time.Duration    // 写超时时间
	reqMemPool    *sync.Pool       // 为每个conn分配一个固定的接收缓存
	rbufMemPool   *sync.Pool       // 为每个conn分配一个环形缓冲区
}

func NewServer() *Server {
	return &Server{
		exitCh:        make(chan bool),
		waitGroup:     &sync.WaitGroup{},
		funcMap:       funcmap.NewFuncMap(),
		acceptTimeout: 30,
		readTimeout:   60,
		writeTimeout:  60,
		reqMemPool: &sync.Pool{
			New: func() interface{} {
				return make([]byte, 1024)
			},
		},
		rbufMemPool: &sync.Pool{
			New: func() interface{} {
				return ringbuffer.NewRingBuffer(1024)
			},
		},
	}
}

func (this *Server) SetAcceptTimeout(acceptTimeout time.Duration) {
	this.acceptTimeout = acceptTimeout
}

func (this *Server) SetReadTimeout(readTimeout time.Duration) {
	this.readTimeout = readTimeout
}

func (this *Server) SetWriteTimeout(writeTimeout time.Duration) {
	this.writeTimeout = writeTimeout
}

func (this *Server) Start(listener *net.TCPListener) {
	log.Printf("Start listen on %v\r\n", listener.Addr())
	this.waitGroup.Add(1)
	defer func() {
		listener.Close()
		this.waitGroup.Done()
	}()

	// 防止恶意连接
	go this.dealSpamConn()

	// report记录，定时发送邮件
	go report.Work()

	for {
		select {
		case <-this.exitCh:
			log.Printf("Stop listen on %v\r\n", listener.Addr())
			return
		default:
		}

		listener.SetDeadline(time.Now().Add(this.acceptTimeout))
		conn, err := listener.AcceptTCP()

		if err != nil {
			if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
				// log.Printf("Accept timeout: %v\r\n", opErr)
				continue
			}
			report.AddCount(report.TryConnect, 1)
			log.Printf("Accept error: %v\r\n", err)
			continue
		}

		report.AddCount(report.SuccessConnect, 1)

		// 连接后等待登陆验证
		handlers.ConnMapLoginStatus.Set(conn, time.Now())
		log.Printf("Accept: %v\r\n", conn.RemoteAddr())

		go this.handleClientConn(conn)
	}
}

// 处理恶意连接，定时检测。
// 若conn的loginstatus为nil,则说明conn已经登陆成功。
// 若conn的loginstatus不为nil,则表示loginstatus为该conn连接服务器时的时间戳(time.Time)
// 判断这个时间戳是否已经超过登陆限制时间，若超过，则断开。
func (this *Server) dealSpamConn() {
	limitTime := 60 * time.Second
	ticker := time.NewTicker(limitTime)
	for _ = range ticker.C {
		items := handlers.ConnMapLoginStatus.Items()
		for conn, loginstatus := range items {
			if loginstatus != nil {
				deadline := loginstatus.(time.Time).Add(limitTime)
				if time.Now().After(deadline) {
					conn.(*net.TCPConn).Close()
					handlers.ConnMapLoginStatus.Delete(conn.(*net.TCPConn))
				}
			}
		}

		// 输出当前的连接数
		fmt.Println("当前连接数: ", handlers.ConnMapUuid.Size())
	}
}

func (this *Server) Stop() {
	// close后，所有的exitCh都返回false
	close(this.exitCh)
	this.waitGroup.Wait()
}

func (this *Server) BindMsgHandler(pacType uint32, fn interface{}) error {
	return this.funcMap.Bind(pacType, fn)
}

func (this *Server) handleClientConn(conn *net.TCPConn) {
	this.waitGroup.Add(1)
	defer this.waitGroup.Done()

	receivePackets := make(chan *packet.Packet, 20) // 接收到的包
	chStop := make(chan bool)                       // 通知停止消息处理
	addr := conn.RemoteAddr().String()

	defer func() {
		defer func() {
			if e := recover(); e != nil {
				log.Printf("Panic: %v\r\n", e)
			}
		}()

		handlers.CloseConn(conn)

		log.Printf("Disconnect: %v\r\n", addr)
		chStop <- true
	}()

	// 处理接收到的包
	go this.handlePackets(conn, receivePackets, chStop)

	// 接收数据
	log.Printf("HandleClient: %v\r\n", addr)

	request := this.reqMemPool.Get().([]byte)
	defer this.reqMemPool.Put(request)

	rbuf := this.rbufMemPool.Get().(*ringbuffer.RingBuffer)
	defer this.rbufMemPool.Put(rbuf)

	for {
		select {
		case <-this.exitCh:
			log.Printf("Stop handleClientConn\r\n")
			return
		default:
		}

		conn.SetReadDeadline(time.Now().Add(this.readTimeout))
		readSize, err := conn.Read(request)
		if err != nil {
			log.Printf("Read failed: %v\r\n", err)
			return
		}

		if readSize > 0 {
			rbuf.Write(request[:readSize])

			// 包长(4) + 类型(4) + 包体(len([]byte))
			for {
				if rbuf.Len() >= 8 {
					pacLen := convert.BytesToUint32(rbuf.Bytes(4))
					if rbuf.Len() >= int(pacLen) {
						rbuf.Peek(4)
						receivePackets <- &packet.Packet{
							Len:  pacLen,
							Type: convert.BytesToUint32(rbuf.Read(4)),
							Data: rbuf.Read(int(pacLen) - 8),
						}
					} else {
						break
					}
				} else {
					break
				}
			}
		}

	}
}

func (this *Server) handlePackets(conn *net.TCPConn, receivePackets <-chan *packet.Packet, chStop <-chan bool) {
	defer func() {
		if e := recover(); e != nil {
			log.Printf("Panic: %v\r\n", e)
		}
	}()

	for {
		select {
		case <-chStop:
			log.Printf("Stop handle receivePackets.\r\n")
			return

		// 消息包处理
		case p := <-receivePackets:
			// 防止模拟的客户端未经登陆就发送消息
			// 如果接收的不是登陆包，则检查是否已经在线，若没在线，则无视消息包，等待登陆检测机制处理。
			if p.Type != packet.PK_ClientLogin {
				if !handlers.ConnMapUuid.Check(conn) {
					continue
				}
			}

			if this.funcMap.Exist(p.Type) {
				this.funcMap.Call(p.Type, conn, p)
			} else {
				log.Printf("Unknown packet type\r\n")
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"

	"github.com/gorilla/handlers"
	"net/http"
)

func main() {
	if httpsErr := BindHTTPS(); httpsErr != nil {
		log.Print(httpsErr)
		if httpErr := BindHTTP(); httpErr != nil {
			log.Print(httpErr)
		}
	} else {
		if httpErr := RedirectHTTP(); httpErr != nil {
			log.Print(httpErr)
		}
	}
	select {}
}

func BindHTTPS() error {
	bindAddress := os.Getenv("FLASHBACK_HTTPS_BIND")
	if len(bindAddress) == 0 {
		return fmt.Errorf("FLASHBACK_HTTPS_BIND notset, not serving HTTPS")
	}
	log.Printf("Serving static assets via HTTPS on %s\n", bindAddress)
	go func() {
		log.Fatal(http.ListenAndServeTLS(bindAddress, os.Getenv("FLASHBACK_SSL_CERT"), os.Getenv("FLASHBACK_SSL_KEY"),
			handlers.LoggingHandler(os.Stderr, http.FileServer(http.Dir("www")))))
	}()
	return nil
}

func BindHTTP() error {
	bindAddress := os.Getenv("FLASHBACK_HTTP_BIND")
	if len(bindAddress) == 0 {
		return fmt.Errorf("FLASHBACK_HTTP_BIND not set, not serving HTTP")
	}
	log.Printf("Redirecting for static assets via HTTP on %s\n\n", bindAddress)
	go func() {
		log.Fatal(http.ListenAndServe(bindAddress, handlers.LoggingHandler(os.Stderr, http.FileServer(http.Dir("www")))))
	}()
	return nil
}

func RedirectHTTP() error {
	bindAddress := os.Getenv("FLASHBACK_HTTP_BIND")
	if len(bindAddress) == 0 {
		return fmt.Errorf("FLASHBACK_HTTP_BIND not set, not redirecting HTTP")
	}
	log.Printf("Serving static assets via HTTP on %s\n", bindAddress)
	go func() {
		log.Fatal(http.ListenAndServe(bindAddress, handlers.LoggingHandler(os.Stderr, http.HandlerFunc(RedirectHandler))))
	}()
	return nil
}

func RedirectHandler(w http.ResponseWriter, r *http.Request) {
	baseURI := os.Getenv("FLASHBACK_BASEURI")
	if len(baseURI) == 0 {
		log.Fatal("FLASHBACK_BASEURI must be set if FLASHBACK_HTTP_BIND is set\n")
	}
	// Redirect all HTTP requests to HTTPS
	// It is the responsibility of the admin to configure FLASHBACK_API_BASEURI
	// properly.  I will not be held responsible for broken redirections
	// or redirect loops!
	http.Redirect(w, r, baseURI, http.StatusFound)
}
<commit_msg>remove old server directory<commit_after><|endoftext|>