<commit_before>package server

import (
	"github.com/patrickmn/go-cache"
	"github.com/pterodactyl/wings/config"
	"github.com/remeh/sizedwaitgroup"
	"go.uber.org/zap"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"
)

// High level definition for a server instance being controlled by Wings.
type Server struct {
	// The unique identifier for the server that should be used when referencing
	// it aganist the Panel API (and internally). This will be used when naming
	// docker containers as well as in log output.
	Uuid string `json:"uuid"`

	// Wether or not the server is in a suspended state. Suspended servers cannot
	// be started or modified except in certain scenarios by an admin user.
	Suspended bool `json:"suspended"`

	// The power state of the server.
	State int `json:"state"`

	// The command that should be used when booting up the server instance.
	Invocation string `json:"invocation"`

	// An array of environment variables that should be passed along to the running
	// server process.
	EnvVars map[string]string `json:"environment" yaml:"env"`

	Build       *BuildSettings `json:"build"`
	Allocations *Allocations   `json:"allocations"`

	Container struct {
		// Defines the Docker image that will be used for this server
		Image string `json:"image,omitempty"`
	} `json:"container,omitempty"`

	environment Environment

	fs *Filesystem

	// Server cache used to store frequently requested information in memory and make
	// certain long operations return faster. For example, FS disk space usage.
	cache *cache.Cache
}

// The build settings for a given server that impact docker container creation and
// resource limits for a server instance.
type BuildSettings struct {
	// The total amount of memory in megabytes that this server is allowed to
	// use on the host system.
	MemoryLimit int64 `json:"memory_limit" yaml:"memory"`

	// The amount of additional swap space to be provided to a container instance.
	Swap int64 `json:"swap"`

	// The relative weight for IO operations in a container. This is relative to other
	// containers on the system and should be a value between 10 and 1000.
	IoWeight uint16 `json:"io_weight" yaml:"io"`

	// The percentage of CPU that this instance is allowed to consume relative to
	// the host. A value of 200% represents complete utilization of two cores. This
	// should be a value between 1 and THREAD_COUNT * 100.
	CpuLimit int64 `json:"cpu_limit" yaml:"cpu"`

	// The amount of disk space in megabytes that a server is allowed to use.
	DiskSpace int64 `json:"disk_space" yaml:"disk"`
}

// Converts the CPU limit for a server build into a number that can be better understood
// by the Docker environment. If there is no limit set, return -1 which will indicate to
// Docker that it has unlimited CPU quota.
func (b *BuildSettings) ConvertedCpuLimit() int64 {
	if b.CpuLimit == 0 {
		return -1
	}

	return b.CpuLimit * 1000
}

// Returns the amount of swap available as a total in bytes. This is returned as the amount
// of memory available to the server initially, PLUS the amount of additional swap to include
// which is the format used by Docker.
func (b *BuildSettings) ConvertedSwap() int64 {
	if b.Swap < 0 {
		return -1
	}

	return (b.Swap * 1000000) + (b.MemoryLimit * 1000000)
}

// Defines the allocations available for a given server. When using the Docker environment
// driver these correspond to mappings for the container that allow external connections.
type Allocations struct {
	// Defines the default allocation that should be used for this server. This is
	// what will be used for {SERVER_IP} and {SERVER_PORT} when modifying configuration
	// files or the startup arguments for a server.
	DefaultMapping struct {
		Ip   string `json:"ip"`
		Port int    `json:"port"`
	} `json:"default" yaml:"default"`

	// Mappings contains all of the ports that should be assigned to a given server
	// attached to the IP they correspond to.
	Mappings map[string][]int `json:"mappings"`
}

// Iterates over a given directory and loads all of the servers listed before returning
// them to the calling function.
func LoadDirectory(dir string, cfg *config.SystemConfiguration) ([]*Server, error) {
	// We could theoretically use a standard wait group here, however doing
	// that introduces the potential to crash the program due to too many
	// open files. This wouldn't happen on a small setup, but once the daemon is
	// handling many servers you run that risk.
	//
	// For now just process 10 files at a time, that should be plenty fast to
	// read and parse the YAML. We should probably make this configurable down
	// the road to help big instances scale better.
	wg := sizedwaitgroup.New(10)

	f, err := ioutil.ReadDir(dir)
	if err != nil {
		return nil, err
	}

	var servers []*Server

	for _, file := range f {
		if !strings.HasSuffix(file.Name(), ".yml") || file.IsDir() {
			continue
		}

		wg.Add()
		// For each of the YAML files we find, parse it and create a new server
		// configuration object that can then be returned to the caller.
		go func(file os.FileInfo) {
			defer wg.Done()

			b, err := ioutil.ReadFile(path.Join(dir, file.Name()))
			if err != nil {
				zap.S().Errorw("failed to read server configuration file, skipping...", zap.Error(err))
				return
			}

			s, err := FromConfiguration(b, cfg)
			if err != nil {
				zap.S().Errorw("failed to parse server configuration, skipping...", zap.Error(err))
				return
			}

			servers = append(servers, s)
		}(file)
	}

	// Wait until we've processed all of the configuration files in the directory
	// before continuing.
	wg.Wait()

	return servers, nil
}

// Initalizes a server using a data byte array. This will be marshaled into the
// given struct using a YAML marshaler. This will also configure the given environment
// for a server.
func FromConfiguration(data []byte, cfg *config.SystemConfiguration) (*Server, error) {
	s := &Server{}

	if err := yaml.Unmarshal(data, s); err != nil {
		return nil, err
	}

	withConfiguration := func (e *DockerEnvironment) {
		e.User = cfg.User.Uid
		e.TimezonePath = cfg.TimezonePath
		e.Server = s
	}

	// Right now we only support a Docker based environment, so I'm going to hard code
	// this logic in. When we're ready to support other environment we'll need to make
	// some modifications here obviously.
	var env Environment
	if t, err := NewDockerEnvironment(withConfiguration); err == nil {
		env = t
	} else {
		return nil, err
	}

	s.environment = env
	s.cache = cache.New(time.Minute * 10, time.Minute * 15)
	s.fs = &Filesystem{
		Root:   cfg.Data,
		Server: s,
	}

	return s, nil
}

// Reads the log file for a server up to a specified number of bytes.
func (s *Server) ReadLogfile(len int64) ([]string, error) {
	return s.Environment().Readlog(len)
}

func (s *Server) Environment() Environment {
	return s.environment
}

func (s *Server) Filesystem() *Filesystem {
	return s.fs
}

func (s *Server) Cache() *cache.Cache {
	return s.cache
}

// Determine if the server is bootable in it's current state or not. This will not
// indicate why a server is not bootable, only if it is.
func (s *Server) IsBootable() bool {
	return s.environment.Exists()
}

// Initalizes a server instance. This will run through and ensure that the environment
// for the server is setup, and that all of the necessary files are created.
func (s *Server) CreateEnvironment() error {
	return s.environment.Create()
}
<commit_msg>Very basic state management for servers<commit_after>package server

import (
	"github.com/patrickmn/go-cache"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/config"
	"github.com/remeh/sizedwaitgroup"
	"go.uber.org/zap"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"
)

// Defines states that identify if the server is running or not.
const (
	StateOffline = "off"
	StateStarting = "starting"
	StateRunning = "running"
	StateStopping = "stopping"
)

type ProcessState string

// High level definition for a server instance being controlled by Wings.
type Server struct {
	// The unique identifier for the server that should be used when referencing
	// it aganist the Panel API (and internally). This will be used when naming
	// docker containers as well as in log output.
	Uuid string `json:"uuid"`

	// Wether or not the server is in a suspended state. Suspended servers cannot
	// be started or modified except in certain scenarios by an admin user.
	Suspended bool `json:"suspended"`

	// The power state of the server.
	State ProcessState `json:"state"`

	// The command that should be used when booting up the server instance.
	Invocation string `json:"invocation"`

	// An array of environment variables that should be passed along to the running
	// server process.
	EnvVars map[string]string `json:"environment" yaml:"env"`

	Build       *BuildSettings `json:"build"`
	Allocations *Allocations   `json:"allocations"`

	Container struct {
		// Defines the Docker image that will be used for this server
		Image string `json:"image,omitempty"`
	} `json:"container,omitempty"`

	environment Environment

	fs *Filesystem

	// Server cache used to store frequently requested information in memory and make
	// certain long operations return faster. For example, FS disk space usage.
	cache *cache.Cache
}

// The build settings for a given server that impact docker container creation and
// resource limits for a server instance.
type BuildSettings struct {
	// The total amount of memory in megabytes that this server is allowed to
	// use on the host system.
	MemoryLimit int64 `json:"memory_limit" yaml:"memory"`

	// The amount of additional swap space to be provided to a container instance.
	Swap int64 `json:"swap"`

	// The relative weight for IO operations in a container. This is relative to other
	// containers on the system and should be a value between 10 and 1000.
	IoWeight uint16 `json:"io_weight" yaml:"io"`

	// The percentage of CPU that this instance is allowed to consume relative to
	// the host. A value of 200% represents complete utilization of two cores. This
	// should be a value between 1 and THREAD_COUNT * 100.
	CpuLimit int64 `json:"cpu_limit" yaml:"cpu"`

	// The amount of disk space in megabytes that a server is allowed to use.
	DiskSpace int64 `json:"disk_space" yaml:"disk"`
}

// Converts the CPU limit for a server build into a number that can be better understood
// by the Docker environment. If there is no limit set, return -1 which will indicate to
// Docker that it has unlimited CPU quota.
func (b *BuildSettings) ConvertedCpuLimit() int64 {
	if b.CpuLimit == 0 {
		return -1
	}

	return b.CpuLimit * 1000
}

// Returns the amount of swap available as a total in bytes. This is returned as the amount
// of memory available to the server initially, PLUS the amount of additional swap to include
// which is the format used by Docker.
func (b *BuildSettings) ConvertedSwap() int64 {
	if b.Swap < 0 {
		return -1
	}

	return (b.Swap * 1000000) + (b.MemoryLimit * 1000000)
}

// Defines the allocations available for a given server. When using the Docker environment
// driver these correspond to mappings for the container that allow external connections.
type Allocations struct {
	// Defines the default allocation that should be used for this server. This is
	// what will be used for {SERVER_IP} and {SERVER_PORT} when modifying configuration
	// files or the startup arguments for a server.
	DefaultMapping struct {
		Ip   string `json:"ip"`
		Port int    `json:"port"`
	} `json:"default" yaml:"default"`

	// Mappings contains all of the ports that should be assigned to a given server
	// attached to the IP they correspond to.
	Mappings map[string][]int `json:"mappings"`
}

// Iterates over a given directory and loads all of the servers listed before returning
// them to the calling function.
func LoadDirectory(dir string, cfg *config.SystemConfiguration) ([]*Server, error) {
	// We could theoretically use a standard wait group here, however doing
	// that introduces the potential to crash the program due to too many
	// open files. This wouldn't happen on a small setup, but once the daemon is
	// handling many servers you run that risk.
	//
	// For now just process 10 files at a time, that should be plenty fast to
	// read and parse the YAML. We should probably make this configurable down
	// the road to help big instances scale better.
	wg := sizedwaitgroup.New(10)

	f, err := ioutil.ReadDir(dir)
	if err != nil {
		return nil, err
	}

	var servers []*Server

	for _, file := range f {
		if !strings.HasSuffix(file.Name(), ".yml") || file.IsDir() {
			continue
		}

		wg.Add()
		// For each of the YAML files we find, parse it and create a new server
		// configuration object that can then be returned to the caller.
		go func(file os.FileInfo) {
			defer wg.Done()

			b, err := ioutil.ReadFile(path.Join(dir, file.Name()))
			if err != nil {
				zap.S().Errorw("failed to read server configuration file, skipping...", zap.Error(err))
				return
			}

			s, err := FromConfiguration(b, cfg)
			if err != nil {
				zap.S().Errorw("failed to parse server configuration, skipping...", zap.Error(err))
				return
			}

			servers = append(servers, s)
		}(file)
	}

	// Wait until we've processed all of the configuration files in the directory
	// before continuing.
	wg.Wait()

	return servers, nil
}

// Initalizes a server using a data byte array. This will be marshaled into the
// given struct using a YAML marshaler. This will also configure the given environment
// for a server.
func FromConfiguration(data []byte, cfg *config.SystemConfiguration) (*Server, error) {
	s := &Server{}

	if err := yaml.Unmarshal(data, s); err != nil {
		return nil, err
	}

	withConfiguration := func (e *DockerEnvironment) {
		e.User = cfg.User.Uid
		e.TimezonePath = cfg.TimezonePath
		e.Server = s
	}

	// Right now we only support a Docker based environment, so I'm going to hard code
	// this logic in. When we're ready to support other environment we'll need to make
	// some modifications here obviously.
	var env Environment
	if t, err := NewDockerEnvironment(withConfiguration); err == nil {
		env = t
	} else {
		return nil, err
	}

	s.environment = env
	s.cache = cache.New(time.Minute * 10, time.Minute * 15)
	s.fs = &Filesystem{
		Root:   cfg.Data,
		Server: s,
	}

	return s, nil
}

// Reads the log file for a server up to a specified number of bytes.
func (s *Server) ReadLogfile(len int64) ([]string, error) {
	return s.Environment().Readlog(len)
}

func (s *Server) Environment() Environment {
	return s.environment
}

func (s *Server) Filesystem() *Filesystem {
	return s.fs
}

func (s *Server) Cache() *cache.Cache {
	return s.cache
}

// Sets the state of the server process.
func (s *Server) SetState(state ProcessState) error {
	switch state {
	case StateOffline:
	case StateRunning:
	case StateStarting:
	case StateStopping:
		s.State = state
		break
	default:
		return errors.New("invalid state provided")
	}

	return nil
}

// Determine if the server is bootable in it's current state or not. This will not
// indicate why a server is not bootable, only if it is.
func (s *Server) IsBootable() bool {
	return s.environment.Exists()
}

// Initalizes a server instance. This will run through and ensure that the environment
// for the server is setup, and that all of the necessary files are created.
func (s *Server) CreateEnvironment() error {
	return s.environment.Create()
}
<|endoftext|>


<commit_before>package server

import (
    "github.com/kataras/iris"
    "github.com/BluePecker/JwtAuth/server/router"
    "fmt"
    "github.com/Sirupsen/logrus"
    //"github.com/kataras/iris/middleware/logger"
    "github.com/kataras/iris/context"
    "strconv"
    "time"
)

type TLS struct {
    Cert string
    Key  string
}

type Options struct {
    Host string
    Port int
    Tls  *TLS
}

type Server struct {
    app *iris.Application
}

func (s *Server) initHttpApp() {
    if s.app == nil {
        s.app = iris.New()
        //s.app.Use(logger.New(logger.Config{
        //    // Status displays status code
        //    Status: true,
        //    // IP displays request's remote address
        //    IP: true,
        //    // Method displays the http method
        //    Method: true,
        //    // Path displays the request path
        //    Path: true,
        //}))

        s.app.Use(func(ctx context.Context) {
            startTime := time.Now()
            ctx.Next()
            endTime := time.Now()
            latency := endTime.Sub(startTime)
            logrus.Infof("%v %4v %s %s %s", strconv.Itoa(ctx.GetStatusCode()), latency, ctx.RemoteAddr(), ctx.Method(), ctx.Path())
        })
    }
}

func (s *Server) Accept(options Options) {
    s.initHttpApp()

    if options.Tls != nil {
        //是否将80端口的请求转发到443
        //target, _ := url.Parse("https://127.0.0.1:443")
        //go host.NewProxy("127.0.0.1:80", target).ListenAndServe()
        var addr string = fmt.Sprintf("%s:%d", options.Host, options.Port)
        err := s.app.Run(iris.TLS(addr, options.Tls.Cert, options.Tls.Key))
        if err != nil {
            logrus.Error(err)
        }

    } else {
        var addr string = fmt.Sprintf("%s:%d", options.Host, options.Port)
        err := s.app.Run(iris.Addr(addr))
        if err != nil {
            logrus.Error(err)
        }
    }
}

func (s *Server) AddRouter(routers... router.Router) {
    s.initHttpApp()

    for _, item := range routers {
        item.Routes(s.app)
    }
}<commit_msg>fix bug<commit_after>package server

import (
    "strconv"
    "time"
    "fmt"
    "github.com/kataras/iris"
    "github.com/BluePecker/JwtAuth/server/router"
    "github.com/Sirupsen/logrus"
    "github.com/kataras/iris/context"
)

type TLS struct {
    Cert string
    Key  string
}

type Options struct {
    Host string
    Port int
    Tls  *TLS
}

type Server struct {
    app *iris.Application
}

func (s *Server) initHttpApp() {
    if s.app == nil {
        s.app = iris.New()
        s.app.Use(func(ctx context.Context) {
            startTime := time.Now()
            ctx.Next()
            endTime := time.Now()
            latency := endTime.Sub(startTime)
            logrus.Infof("%v %4v %s %s %s", strconv.Itoa(ctx.GetStatusCode()), latency, ctx.RemoteAddr(), ctx.Method(), ctx.Path())
        })
    }
}

func (s *Server) Accept(options Options) {
    s.initHttpApp()

    if options.Tls != nil {
        //是否将80端口的请求转发到443
        //target, _ := url.Parse("https://127.0.0.1:443")
        //go host.NewProxy("127.0.0.1:80", target).ListenAndServe()
        var addr string = fmt.Sprintf("%s:%d", options.Host, options.Port)
        err := s.app.Run(iris.TLS(addr, options.Tls.Cert, options.Tls.Key))
        if err != nil {
            logrus.Error(err)
        }

    } else {
        var addr string = fmt.Sprintf("%s:%d", options.Host, options.Port)
        err := s.app.Run(iris.Addr(addr))
        if err != nil {
            logrus.Error(err)
        }
    }
}

func (s *Server) AddRouter(routers... router.Router) {
    s.initHttpApp()

    for _, item := range routers {
        item.Routes(s.app)
    }
}<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/seiflotfy/skizze/config"
	"github.com/seiflotfy/skizze/counters"
	"github.com/seiflotfy/skizze/storage"
	"github.com/seiflotfy/skizze/utils"
)

type requestData struct {
	id       string
	typ      string
	Capacity uint64   `json:"capacity"`
	Values   []string `json:"values"`
}

var logger = utils.GetLogger()
var counterManager *counters.ManagerStruct

/*
Server manages the http connections and communciates with the counters manager
*/
type Server struct{}

type sketchesResult struct {
	Result []string `json:"result"`
	Error  error    `json:"error"`
}

type sketchResult struct {
	Result interface{} `json:"result"`
	Error  error       `json:"error"`
}

/*
New returns a new Server
*/
func New() (*Server, error) {
	var err error
	counterManager, err = counters.GetManager()
	if err != nil {
		return nil, err
	}
	server := Server{}
	return &server, nil
}

func (srv *Server) handleTopRequest(w http.ResponseWriter, method string, data requestData) {
	var err error
	var sketches []string
	var js []byte

	switch {
	case method == "GET":
		// Get all counters
		sketches, err = counterManager.GetSketchs()
		js, err = json.Marshal(sketchesResult{sketches, err})
		logger.Info.Printf("[%v]: Getting all available sketches", method)
	case method == "MERGE":
		// Reserved for merging hyper log log
		http.Error(w, "Not Implemented", http.StatusNotImplemented)
		return
	default:
		http.Error(w, "Invalid Method: "+method, http.StatusBadRequest)
		return
	}

	if err == nil {
		w.Header().Set("Content-Type", "application/json")
		w.Write(js)
	} else {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

}

func (srv *Server) handleSketchRequest(w http.ResponseWriter, method string, data requestData) {
	var res sketchResult
	var err error

	// TODO (mb): handle errors from counterManager.*
	switch {
	case method == "GET":
		// Get a count for a specific sketch
		count, err := counterManager.GetCountForSketch(data.id, data.Values)
		logger.Info.Printf("[%v]: Getting counter for sketch: %v", method, data.id)
		res = sketchResult{count, err}
	case method == "POST":
		// Create a new sketch counter
		err = counterManager.CreateSketch(data.id, data.typ, data.Capacity)
		logger.Info.Printf("[%v]: Creating new sketch: %v", method, data.id)
		res = sketchResult{0, err}
	case method == "PUT":
		// Add values to counter
		err = counterManager.AddToSketch(data.id, data.Values)
		logger.Info.Printf("[%v]: Updating counter for sketch: %v", method, data.id)
		res = sketchResult{nil, err}
	case method == "PURGE":
		// Purges values from counter
		err = counterManager.DeleteFromSketch(data.id, data.Values)
		logger.Info.Printf("[%v]: Purging values for sketch: %v", method, data.id)
		res = sketchResult{nil, err}
	case method == "DELETE":
		// Delete Counter
		err := counterManager.DeleteSketch(data.id)
		logger.Info.Printf("[%v]: Deleting sketch: %v", method, data.id)
		res = sketchResult{nil, err}
	default:
		logger.Error.Printf("[%v]: Invalid Method: %v", method, http.StatusBadRequest)
		http.Error(w, fmt.Sprintf("Invalid Method: %s", method), http.StatusBadRequest)
		return
	}

	js, err := json.Marshal(res)

	if err == nil {
		w.Header().Set("Content-Type", "application/json")
		w.Write(js)
	} else {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (srv *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	method := r.Method
	paths := strings.Split(r.URL.Path[1:], "/")
	body, _ := ioutil.ReadAll(r.Body)
	var data requestData
	if len(body) > 0 {
		err := json.Unmarshal(body, &data)
		if err != nil {
			logger.Error.Printf("An error has ocurred: %v", err.Error())
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
	} else {
		data = requestData{}
	}

	if len(paths) == 1 {
		srv.handleTopRequest(w, method, data)
	} else if len(paths) == 2 {
		data.typ = strings.TrimSpace(string(paths[0]))
		data.id = strings.TrimSpace(strings.Join(paths[1:], "/"))
		srv.handleSketchRequest(w, method, data)
	}
}

/*
Run ...
*/
func (srv *Server) Run() {
	conf := config.GetConfig()
	port := int(conf.GetPort())
	logger.Info.Println("Server up and running on port: " + strconv.Itoa(port))
	http.ListenAndServe(":"+strconv.Itoa(port), srv)
}

/*
Stop ...
*/
func (srv *Server) Stop() {
	logger.Info.Println("Stopping server...")
	storage.CloseInfoDB()
	os.Exit(0)
}
<commit_msg>Fix bug with wrong method call<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/seiflotfy/skizze/config"
	"github.com/seiflotfy/skizze/counters"
	"github.com/seiflotfy/skizze/storage"
	"github.com/seiflotfy/skizze/utils"
)

type requestData struct {
	id       string
	typ      string
	Capacity uint64   `json:"capacity"`
	Values   []string `json:"values"`
}

var logger = utils.GetLogger()
var counterManager *counters.ManagerStruct

/*
Server manages the http connections and communciates with the counters manager
*/
type Server struct{}

type sketchesResult struct {
	Result []string `json:"result"`
	Error  error    `json:"error"`
}

type sketchResult struct {
	Result interface{} `json:"result"`
	Error  error       `json:"error"`
}

/*
New returns a new Server
*/
func New() (*Server, error) {
	var err error
	counterManager, err = counters.GetManager()
	if err != nil {
		return nil, err
	}
	server := Server{}
	return &server, nil
}

func (srv *Server) handleTopRequest(w http.ResponseWriter, method string, data requestData) {
	var err error
	var sketches []string
	var js []byte

	switch {
	case method == "GET":
		// Get all counters
		sketches, err = counterManager.GetSketches()
		js, err = json.Marshal(sketchesResult{sketches, err})
		logger.Info.Printf("[%v]: Getting all available sketches", method)
	case method == "MERGE":
		// Reserved for merging hyper log log
		http.Error(w, "Not Implemented", http.StatusNotImplemented)
		return
	default:
		http.Error(w, "Invalid Method: "+method, http.StatusBadRequest)
		return
	}

	if err == nil {
		w.Header().Set("Content-Type", "application/json")
		w.Write(js)
	} else {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

}

func (srv *Server) handleSketchRequest(w http.ResponseWriter, method string, data requestData) {
	var res sketchResult
	var err error

	// TODO (mb): handle errors from counterManager.*
	switch {
	case method == "GET":
		// Get a count for a specific sketch
		count, err := counterManager.GetCountForSketch(data.id, data.Values)
		logger.Info.Printf("[%v]: Getting counter for sketch: %v", method, data.id)
		res = sketchResult{count, err}
	case method == "POST":
		// Create a new sketch counter
		err = counterManager.CreateSketch(data.id, data.typ, data.Capacity)
		logger.Info.Printf("[%v]: Creating new sketch: %v", method, data.id)
		res = sketchResult{0, err}
	case method == "PUT":
		// Add values to counter
		err = counterManager.AddToSketch(data.id, data.Values)
		logger.Info.Printf("[%v]: Updating counter for sketch: %v", method, data.id)
		res = sketchResult{nil, err}
	case method == "PURGE":
		// Purges values from counter
		err = counterManager.DeleteFromSketch(data.id, data.Values)
		logger.Info.Printf("[%v]: Purging values for sketch: %v", method, data.id)
		res = sketchResult{nil, err}
	case method == "DELETE":
		// Delete Counter
		err := counterManager.DeleteSketch(data.id)
		logger.Info.Printf("[%v]: Deleting sketch: %v", method, data.id)
		res = sketchResult{nil, err}
	default:
		logger.Error.Printf("[%v]: Invalid Method: %v", method, http.StatusBadRequest)
		http.Error(w, fmt.Sprintf("Invalid Method: %s", method), http.StatusBadRequest)
		return
	}

	js, err := json.Marshal(res)

	if err == nil {
		w.Header().Set("Content-Type", "application/json")
		w.Write(js)
	} else {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (srv *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	method := r.Method
	paths := strings.Split(r.URL.Path[1:], "/")
	body, _ := ioutil.ReadAll(r.Body)
	var data requestData
	if len(body) > 0 {
		err := json.Unmarshal(body, &data)
		if err != nil {
			logger.Error.Printf("An error has ocurred: %v", err.Error())
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
	} else {
		data = requestData{}
	}

	if len(paths) == 1 {
		srv.handleTopRequest(w, method, data)
	} else if len(paths) == 2 {
		data.typ = strings.TrimSpace(string(paths[0]))
		data.id = strings.TrimSpace(strings.Join(paths[1:], "/"))
		srv.handleSketchRequest(w, method, data)
	}
}

/*
Run ...
*/
func (srv *Server) Run() {
	conf := config.GetConfig()
	port := int(conf.GetPort())
	logger.Info.Println("Server up and running on port: " + strconv.Itoa(port))
	http.ListenAndServe(":"+strconv.Itoa(port), srv)
}

/*
Stop ...
*/
func (srv *Server) Stop() {
	logger.Info.Println("Stopping server...")
	storage.CloseInfoDB()
	os.Exit(0)
}
<|endoftext|>


<commit_before>package server

import (
	"net"
	"net/http"
	"sync"

	"github.com/getlantern/golog"
)

var (
	log = golog.LoggerFor("server")
)

type Server struct {
	httpServer http.Server
	tls        bool

	listener net.Listener
	handler  http.Handler
}

func NewServer(handler http.Handler) *Server {
	server := &Server{}
	return server
}

func (s *Server) ServeHTTP(addr string, chListenOn *chan string) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	s.tls = false
	log.Debugf("Listen http on %s", addr)
	return s.doServe(listener, chListenOn)
}

func (s *Server) ServeHTTPS(addr, keyfile, certfile string, chListenOn *chan string) error {
	listener, err := listenTLS(addr, keyfile, certfile)
	if err != nil {
		return err
	}
	s.tls = true
	log.Debugf("Listen https on %s", addr)
	return s.doServe(listener, chListenOn)
}

// connBag is a just bag of connections. You can put a connection in and
// withdraw it afterwards, or purge it regardless it's withdrawed or not.
type connBag struct {
	mu sync.Mutex
	m  map[string]net.Conn
}

func (cb *connBag) Put(c net.Conn) {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	cb.m[c.RemoteAddr().String()] = c
}

func (cb *connBag) Withdraw(remoteAddr string) (c net.Conn) {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	c = cb.m[remoteAddr]
	delete(cb.m, remoteAddr)
	return
}

func (cb *connBag) Purge(remoteAddr string) {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	// non-op if item doesn't exist
	delete(cb.m, remoteAddr)
}

func (s *Server) doServe(listener net.Listener, chListenOn *chan string) error {
	//cb := connBag{m: make(map[string]net.Conn)}

	proxy := http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			/*
				c := cb.Withdraw(req.RemoteAddr)
				context.Set(req, "conn", c)
			*/
			s.handler.ServeHTTP(w, req)
		})

	s.listener = listener

	/*
		limListener := newLimitedListener(listener, &s.numConns, s.idleTimeout)
		preListener := preprocessor.NewListener(limListener)

		if s.enableReports {
			mListener := measured.Listener(preListener, 30*time.Second)
			s.listener = mListener
		} else {
			s.listener = preListener
		}
	*/
	s.httpServer = http.Server{Handler: proxy} /*
		ConnState: func(c net.Conn, state http.ConnState) {
			if sc, ok := c.(preprocessor.StatefulConn); ok {
				sc.SetState(state)
			}
			switch state {
			case http.StateNew:
				if atomic.LoadUint64(&s.numConns) >= s.maxConns {
					limListener.Stop()
				} else if limListener.IsStopped() {
					limListener.Restart()
				}
			case http.StateActive:
				cb.Put(c)
			case http.StateClosed:
				// When go server encounters abnormal request, it
				// will transit to StateClosed directly without
				// the handler being invoked, hence the connection
				// will not be withdrawed. Purge it in such case.
				cb.Purge(c.RemoteAddr().String())
			}
		},
	*/

	/*
		addr := s.listener.Addr().String()
		host, port, err := net.SplitHostPort(addr)
		if err != nil {
			panic("should not happen")
		}
		mimic.Host = host
		mimic.Port = port
		if chListenOn != nil {
			*chListenOn <- addr
		}
	*/

	return s.httpServer.Serve(s.listener)
}
<commit_msg>Added connection bag back<commit_after>package server

import (
	"net"
	"net/http"
	"sync"

	"github.com/gorilla/context"

	"github.com/getlantern/golog"
)

var (
	log = golog.LoggerFor("server")
)

type Server struct {
	httpServer http.Server
	tls        bool

	listener net.Listener
	handler  http.Handler
}

func NewServer(handler http.Handler) *Server {
	server := &Server{}
	return server
}

func (s *Server) ServeHTTP(addr string, chListenOn *chan string) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	s.tls = false
	log.Debugf("Listen http on %s", addr)
	return s.doServe(listener, chListenOn)
}

func (s *Server) ServeHTTPS(addr, keyfile, certfile string, chListenOn *chan string) error {
	listener, err := listenTLS(addr, keyfile, certfile)
	if err != nil {
		return err
	}
	s.tls = true
	log.Debugf("Listen https on %s", addr)
	return s.doServe(listener, chListenOn)
}

// connBag is a just bag of connections. You can put a connection in and
// withdraw it afterwards, or purge it regardless it's withdrawed or not.
type connBag struct {
	mu sync.Mutex
	m  map[string]net.Conn
}

func (cb *connBag) Put(c net.Conn) {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	cb.m[c.RemoteAddr().String()] = c
}

func (cb *connBag) Withdraw(remoteAddr string) (c net.Conn) {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	c = cb.m[remoteAddr]
	delete(cb.m, remoteAddr)
	return
}

func (cb *connBag) Purge(remoteAddr string) {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	// non-op if item doesn't exist
	delete(cb.m, remoteAddr)
}

func (s *Server) doServe(listener net.Listener, chListenOn *chan string) error {
	cb := connBag{m: make(map[string]net.Conn)}

	proxy := http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			c := cb.Withdraw(req.RemoteAddr)
			context.Set(req, "conn", c)
			s.handler.ServeHTTP(w, req)
		})

	s.listener = listener
	/*
		limListener := newLimitedListener(listener, &s.numConns, time.Duration(30)*time.Second)
		preListener := preprocessor.NewListener(limListener)

		if s.enableReports {
			mListener := measured.Listener(preListener, 30*time.Second)
			s.listener = mListener
		} else {
			s.listener = preListener
		}
	*/
	s.httpServer = http.Server{
		Handler: proxy,
		ConnState: func(c net.Conn, state http.ConnState) {
			switch state {
			case http.StateNew:
				/*
					if atomic.LoadUint64(&s.numConns) >= s.maxConns {
						limListener.Stop()
					} else if limListener.IsStopped() {
						limListener.Restart()
					}
				*/
			case http.StateActive:
				cb.Put(c)
			case http.StateClosed:
				// When go server encounters abnormal request, it
				// will transit to StateClosed directly without
				// the handler being invoked, hence the connection
				// will not be withdrawed. Purge it in such case.
				cb.Purge(c.RemoteAddr().String())
			}
		},
	}

	addr := s.listener.Addr().String()
	/*
		host, port, err := net.SplitHostPort(addr)
		if err != nil {
			panic("should not happen")
		}
	*/

	if chListenOn != nil {
		*chListenOn <- addr
	}

	return s.httpServer.Serve(s.listener)
}
<|endoftext|>


<commit_before>package server

import (
	"net"
	"os"
	"strings"

	"github.com/kataras/iris/config"
	"github.com/valyala/fasthttp"
)

// Server is the IServer's implementation, holds the fasthttp's Server, a net.Listener, the ServerOptions, and the handler
// handler is registed at the Station/Iris level
type Server struct {
	*fasthttp.Server
	listener net.Listener
	Config   config.Server
	started  bool
	tls      bool
	handler  fasthttp.RequestHandler
}

// New returns a pointer to a Server object, and set it's options if any,  nothing more
func New(cfg ...config.Server) *Server {
	c := config.DefaultServer().Merge(cfg)
	s := &Server{Server: &fasthttp.Server{Name: config.ServerName}, Config: c}

	s.Config.ListeningAddr = parseAddr(s.Config.ListeningAddr)

	return s
}

// SetHandler sets the handler in order to listen on new requests, this is done at the Station/Iris level
func (s *Server) SetHandler(h fasthttp.RequestHandler) {
	s.handler = h
	if s.Server != nil {
		s.Server.Handler = s.handler
	}
}

// Handler returns the fasthttp.RequestHandler which is registed to the Server
func (s *Server) Handler() fasthttp.RequestHandler {
	return s.handler
}

// IsListening returns true if server is listening/started, otherwise false
func (s *Server) IsListening() bool {
	return s.started
}

// IsSecure returns true if server uses TLS, otherwise false
func (s *Server) IsSecure() bool {
	return s.tls
}

// Listener returns the net.Listener which this server (is) listening to
func (s *Server) Listener() net.Listener {
	return s.listener
}

// Host returns the Listener().Addr().String(), if server is not listening it returns the config.ListeningAddr
func (s *Server) Host() (host string) {
	if s.IsListening() {
		return s.Listener().Addr().String()
	} else {
		return s.Config.ListeningAddr
	}
}

// Hostname returns the hostname part only, if host == 0.0.0.0:8080 it will return the 0.0.0.0
// if server is not listening it returns the config.ListeningAddr's hostname part
func (s *Server) Hostname() (hostname string) {
	if s.IsListening() {
		fullhost := s.Listener().Addr().String()
		hostname = fullhost[0:strings.IndexByte(fullhost, ':')] // no the port
	} else {
		fullhost := s.Config.ListeningAddr
		if idx := strings.IndexByte(fullhost, ':'); idx > 1 { // at least after second char
			hostname = hostname[0 : idx-1]
		} else {
			hostname = "0.0.0.0"
		}

	}
	return
}

//Serve just serves a listener, it is a blocking action, plugin.PostListen is not fired here.
func (s *Server) Serve(l net.Listener) error {
	s.listener = l
	return s.Server.Serve(l)
}

// listen starts the process of listening to the new requests
func (s *Server) listen() (err error) {

	if s.started {
		err = ErrServerAlreadyStarted.Return()
		return
	}
	s.listener, err = net.Listen("tcp4", s.Config.ListeningAddr)

	if err != nil {
		err = ErrServerPortAlreadyUsed.Return()
		return
	}

	//Non-block way here because I want the plugin's PostListen ability...
	go s.Server.Serve(s.listener)

	s.started = true
	s.tls = false

	return
}

// listenTLS starts the process of listening to the new requests using TLS, keyfile and certfile are given before this method fires
func (s *Server) listenTLS() (err error) {

	if s.started {
		err = ErrServerAlreadyStarted.Return()
		return
	}

	if s.Config.CertFile == "" || s.Config.KeyFile == "" {
		err = ErrServerTLSOptionsMissing.Return()
		return
	}

	s.listener, err = net.Listen("tcp4", s.Config.ListeningAddr)

	if err != nil {
		err = ErrServerPortAlreadyUsed.Return()
		return
	}

	go s.Server.ServeTLS(s.listener, s.Config.CertFile, s.Config.KeyFile)

	s.started = true
	s.tls = true

	return
}

// listenUnix  starts the process of listening to the new requests using a 'socket file', this works only on unix
func (s *Server) listenUnix() (err error) {

	if s.started {
		err = ErrServerAlreadyStarted.Return()
		return
	}

	mode := s.Config.Mode

	//this code is from fasthttp ListenAndServeUNIX, I extracted it because we need the tcp.Listener
	if errOs := os.Remove(s.Config.ListeningAddr); errOs != nil && !os.IsNotExist(errOs) {
		err = ErrServerRemoveUnix.Format(s.Config.ListeningAddr, errOs.Error())
		return
	}
	s.listener, err = net.Listen("unix", s.Config.ListeningAddr)

	if err != nil {
		err = ErrServerPortAlreadyUsed.Return()
		return
	}

	if err = os.Chmod(s.Config.ListeningAddr, mode); err != nil {
		err = ErrServerChmod.Format(mode, s.Config.ListeningAddr, err.Error())
		return
	}

	s.Server.Handler = s.handler
	go s.Server.Serve(s.listener)

	s.started = true
	s.tls = false

	return

}

// OpenServer opens/starts/runs/listens (to) the server, listenTLS if Cert && Key is registed, listenUnix if Mode is registed, otherwise listen
// instead of return an error this is panics on any server's error
func (s *Server) OpenServer() (err error) {
	if s.Config.CertFile != "" && s.Config.KeyFile != "" {
		err = s.listenTLS()
	} else if s.Config.Mode > 0 {
		err = s.listenUnix()
	} else {
		err = s.listen()
	}

	return
}

// CloseServer closes the server
func (s *Server) CloseServer() error {

	if !s.started {
		return ErrServerIsClosed.Return()
	}

	if s.listener != nil {
		return s.listener.Close()
	}
	return nil
}

// parseAddr gets a slice of string and returns the address of which the Iris' server can listen
func parseAddr(fullHostOrPort ...string) string {

	if len(fullHostOrPort) > 1 {
		fullHostOrPort = fullHostOrPort[0:1]
	}
	addr := config.DefaultServerAddr // default address
	// if nothing passed, then use environment's port (if any) or just :8080
	if len(fullHostOrPort) == 0 {
		if envPort := os.Getenv("PORT"); len(envPort) > 0 {
			addr = ":" + envPort
		}

	} else if len(fullHostOrPort) == 1 {
		addr = fullHostOrPort[0]
		if strings.IndexRune(addr, ':') == -1 {
			//: doesn't found on the given address, so maybe it's only a port
			addr = ":" + addr
		}
	}

	return addr
}
<commit_msg>A small fix for offline server tests<commit_after>package server

import (
	"net"
	"os"
	"strings"

	"github.com/kataras/iris/config"
	"github.com/valyala/fasthttp"
)

// Server is the IServer's implementation, holds the fasthttp's Server, a net.Listener, the ServerOptions, and the handler
// handler is registed at the Station/Iris level
type Server struct {
	*fasthttp.Server
	listener net.Listener
	Config   config.Server
	started  bool
	tls      bool
	handler  fasthttp.RequestHandler
}

// New returns a pointer to a Server object, and set it's options if any,  nothing more
func New(cfg ...config.Server) *Server {
	c := config.DefaultServer().Merge(cfg)
	s := &Server{Server: &fasthttp.Server{Name: config.ServerName}, Config: c}

	s.Config.ListeningAddr = parseAddr(s.Config.ListeningAddr)

	return s
}

// SetHandler sets the handler in order to listen on new requests, this is done at the Station/Iris level
func (s *Server) SetHandler(h fasthttp.RequestHandler) {
	s.handler = h
	if s.Server != nil {
		s.Server.Handler = s.handler
	}
}

// Handler returns the fasthttp.RequestHandler which is registed to the Server
func (s *Server) Handler() fasthttp.RequestHandler {
	return s.handler
}

// IsListening returns true if server is listening/started, otherwise false
func (s *Server) IsListening() bool {
	return s.started
}

// IsSecure returns true if server uses TLS, otherwise false
func (s *Server) IsSecure() bool {
	return s.tls
}

// Listener returns the net.Listener which this server (is) listening to
func (s *Server) Listener() net.Listener {
	return s.listener
}

// Host returns the Listener().Addr().String(), if server is not listening it returns the config.ListeningAddr
func (s *Server) Host() (host string) {
	if s.IsListening() {
		return s.Listener().Addr().String()
	} else {
		return s.Config.ListeningAddr
	}
}

// Hostname returns the hostname part only, if host == 0.0.0.0:8080 it will return the 0.0.0.0
// if server is not listening it returns the config.ListeningAddr's hostname part
func (s *Server) Hostname() (hostname string) {
	if s.IsListening() {
		fullhost := s.Listener().Addr().String()
		hostname = fullhost[0:strings.IndexByte(fullhost, ':')] // no the port
	} else {
		fullhost := s.Config.ListeningAddr
		if idx := strings.IndexByte(fullhost, ':'); idx > 1 { // at least after second char
			hostname = hostname[0:idx]
		} else {
			hostname = "0.0.0.0"
		}

	}
	return
}

//Serve just serves a listener, it is a blocking action, plugin.PostListen is not fired here.
func (s *Server) Serve(l net.Listener) error {
	s.listener = l
	return s.Server.Serve(l)
}

// listen starts the process of listening to the new requests
func (s *Server) listen() (err error) {

	if s.started {
		err = ErrServerAlreadyStarted.Return()
		return
	}
	s.listener, err = net.Listen("tcp4", s.Config.ListeningAddr)

	if err != nil {
		err = ErrServerPortAlreadyUsed.Return()
		return
	}

	//Non-block way here because I want the plugin's PostListen ability...
	go s.Server.Serve(s.listener)

	s.started = true
	s.tls = false

	return
}

// listenTLS starts the process of listening to the new requests using TLS, keyfile and certfile are given before this method fires
func (s *Server) listenTLS() (err error) {

	if s.started {
		err = ErrServerAlreadyStarted.Return()
		return
	}

	if s.Config.CertFile == "" || s.Config.KeyFile == "" {
		err = ErrServerTLSOptionsMissing.Return()
		return
	}

	s.listener, err = net.Listen("tcp4", s.Config.ListeningAddr)

	if err != nil {
		err = ErrServerPortAlreadyUsed.Return()
		return
	}

	go s.Server.ServeTLS(s.listener, s.Config.CertFile, s.Config.KeyFile)

	s.started = true
	s.tls = true

	return
}

// listenUnix  starts the process of listening to the new requests using a 'socket file', this works only on unix
func (s *Server) listenUnix() (err error) {

	if s.started {
		err = ErrServerAlreadyStarted.Return()
		return
	}

	mode := s.Config.Mode

	//this code is from fasthttp ListenAndServeUNIX, I extracted it because we need the tcp.Listener
	if errOs := os.Remove(s.Config.ListeningAddr); errOs != nil && !os.IsNotExist(errOs) {
		err = ErrServerRemoveUnix.Format(s.Config.ListeningAddr, errOs.Error())
		return
	}
	s.listener, err = net.Listen("unix", s.Config.ListeningAddr)

	if err != nil {
		err = ErrServerPortAlreadyUsed.Return()
		return
	}

	if err = os.Chmod(s.Config.ListeningAddr, mode); err != nil {
		err = ErrServerChmod.Format(mode, s.Config.ListeningAddr, err.Error())
		return
	}

	s.Server.Handler = s.handler
	go s.Server.Serve(s.listener)

	s.started = true
	s.tls = false

	return

}

// OpenServer opens/starts/runs/listens (to) the server, listenTLS if Cert && Key is registed, listenUnix if Mode is registed, otherwise listen
// instead of return an error this is panics on any server's error
func (s *Server) OpenServer() (err error) {
	if s.Config.CertFile != "" && s.Config.KeyFile != "" {
		err = s.listenTLS()
	} else if s.Config.Mode > 0 {
		err = s.listenUnix()
	} else {
		err = s.listen()
	}

	return
}

// CloseServer closes the server
func (s *Server) CloseServer() error {

	if !s.started {
		return ErrServerIsClosed.Return()
	}

	if s.listener != nil {
		return s.listener.Close()
	}
	return nil
}

// parseAddr gets a slice of string and returns the address of which the Iris' server can listen
func parseAddr(fullHostOrPort ...string) string {

	if len(fullHostOrPort) > 1 {
		fullHostOrPort = fullHostOrPort[0:1]
	}
	addr := config.DefaultServerAddr // default address
	// if nothing passed, then use environment's port (if any) or just :8080
	if len(fullHostOrPort) == 0 {
		if envPort := os.Getenv("PORT"); len(envPort) > 0 {
			addr = ":" + envPort
		}

	} else if len(fullHostOrPort) == 1 {
		addr = fullHostOrPort[0]
		if strings.IndexRune(addr, ':') == -1 {
			//: doesn't found on the given address, so maybe it's only a port
			addr = ":" + addr
		}
	}

	return addr
}
<|endoftext|>


<commit_before>package server

import (
    "net/http"
    "github.com/gorilla/websocket"
    "log"
    "fmt"
)

type Server struct {
    upgrader        websocket.Upgrader
    connections     []*websocket.Conn
}

func webSocketHandler(conn *websocket.Conn) {
    log.Println("New websocket connection available")
    for {
        msgType, msg, err := conn.ReadMessage()
        if err != nil {
            log.Println("Connection read error", err.Error())
            return
        }

        if msgType == websocket.TextMessage {
            log.Println(string(msg))
        }

        if err = conn.WriteMessage(msgType, msg); err != nil {
            log.Println("Connection write error", err.Error())
            return
        }
    }
}

func New() *Server {
	return &Server{
    }
}

func (s *Server) Start(host string) {

    http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        fmt.Fprintf(w, "Hello World")
    })

    var upgrader = websocket.Upgrader{
        ReadBufferSize: 1024,
        WriteBufferSize: 1024,
    }

    http.HandleFunc("/websocket", func(w http.ResponseWriter, r *http.Request) {
        conn, err := upgrader.Upgrade(w, r, nil)
        if err != nil {
            log.Println(err)
            return
        }

        go webSocketHandler(conn)
    })

    log.Fatal(http.ListenAndServe(":9000", nil))
}
<commit_msg>Allow all request origins for server connections<commit_after>package server

import (
    "net/http"
    "github.com/gorilla/websocket"
    "log"
    "fmt"
)

type Server struct {
    upgrader        websocket.Upgrader
    connections     []*websocket.Conn
}

func webSocketHandler(conn *websocket.Conn) {
    log.Println("New websocket connection available")
    for {
        msgType, msg, err := conn.ReadMessage()
        if err != nil {
            log.Println("Connection read error", err.Error())
            return
        }

        if msgType == websocket.TextMessage {
            log.Println(string(msg))
        }

        if err = conn.WriteMessage(msgType, msg); err != nil {
            log.Println("Connection write error", err.Error())
            return
        }
    }
}

func New() *Server {
	return &Server{
    }
}

func (s *Server) Start(host string) {

    http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        fmt.Fprintf(w, "Hello World")
    })

    var upgrader = websocket.Upgrader{
        ReadBufferSize: 1024,
        WriteBufferSize: 1024,
        CheckOrigin: func(r *http.Request) bool { return true },
    }

    http.HandleFunc("/websocket", func(w http.ResponseWriter, r *http.Request) {
        conn, err := upgrader.Upgrade(w, r, nil)
        if err != nil {
            log.Println(err)
            return
        }

        go webSocketHandler(conn)
    })

    log.Fatal(http.ListenAndServe(":9000", nil))
}
<|endoftext|>


<commit_before>// Copyright 2016
// 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 (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"time"

	"context"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/manifest/schema2"
	"github.com/docker/distribution/notifications"
	"github.com/mailgun/manners"
	"github.com/nhurel/dim/lib"
	"github.com/nhurel/dim/lib/environment"
)

// Server type handle  indexation of a docker registry and serves the search endpoint
type Server struct {
	*manners.GracefulServer
	index dim.RegistryIndex
}

// NewServer creates a new Server instance to listen on given port and use given index
func NewServer(port string, index dim.RegistryIndex, ctx context.Context) *Server {
	c := environment.Set(ctx, environment.StartTimeKey, time.Now())
	http.HandleFunc("/v1/search", handler(index, Search))
	http.HandleFunc("/dim/notify", handler(index, NotifyImageChange))
	http.HandleFunc("/dim/version", buildVersionHandler(c))
	http.HandleFunc("/v2/_catalog", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, "{}")
	})
	return &Server{manners.NewWithServer(&http.Server{Addr: port, Handler: http.DefaultServeMux}), index}
}

// Run starts the server instance
func (s *Server) Run() error {
	return s.ListenAndServe()
}

// Handler injects an index into an HandlerFunc
func handler(i dim.RegistryIndex, dhf DimHandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		dhf(i, w, r)
	}
}

func buildVersionHandler(ctx context.Context) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		Version(ctx, w, r)
	}
}

// Version return server info including info and uptime
func Version(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	info := dim.Info{}
	v := environment.Get(ctx, environment.VersionKey)
	if v != nil {
		info.Version = v.(string)
	}
	start := environment.Get(ctx, environment.StartTimeKey)
	if start != nil {
		info.Uptime = time.Now().Sub(start.(time.Time)).String()
	}
	if b, err := json.Marshal(info); err != nil {
		http.Error(w, "Failed to serialize the response", http.StatusInternalServerError)
		logrus.WithError(err).Errorln("Error occured while serializing server info")
	} else {
		fmt.Fprint(w, string(b))
	}
}

// NotifyImageChange handles docker registry events
func NotifyImageChange(i dim.RegistryIndex, w http.ResponseWriter, r *http.Request) {

	logrus.Debugln("Receiving event from registry")
	defer r.Body.Close()

	enveloppe := &notifications.Envelope{}

	if err := json.NewDecoder(r.Body).Decode(enveloppe); err != nil {
		logrus.WithError(err).Errorln("Failed to parse event")
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
	fmt.Fprintln(w, "Notiifcation handled !")

	logrus.WithField("enveloppe", enveloppe).Debugln("Processing event")

	for _, event := range enveloppe.Events {
		switch event.Action {
		case notifications.EventActionDelete:
			logrus.WithField("enveloppe", enveloppe).Infoln("Processing delete event")
			i.Submit(&dim.NotificationJob{Action: dim.DeleteAction, Digest: event.Target.Digest})
		// i.DeleteImage(event.Target.Digest.String())
		case notifications.EventActionPush:
			if event.Target.MediaType == schema2.MediaTypeManifest {
				logrus.WithField("enveloppe", enveloppe).Infoln("Processing push event")
				i.Submit(&dim.NotificationJob{Action: dim.PushAction, Repository: event.Target.Repository, Tag: event.Target.Tag, Digest: event.Target.Digest})
				//if err := i.GetImageAndIndex(event.Target.Repository, event.Target.Tag, event.Target.Digest); err != nil {
				//	logrus.WithField("EventTarget", event.Target).WithError(err).Errorln("Failed to reindex image")
				//}
			} else {
				logrus.WithField("mediatype", event.Target.MediaType).WithField("Event", event).Debugln("Event safely ignored because mediatype is unknown")
			}
		default:
			logrus.WithField("Action", event.Action).WithField("Event", event).Debugln("Event safely ignored")
		}
	}
}

// Search handles docker search request
func Search(i dim.RegistryIndex, w http.ResponseWriter, r *http.Request) {

	var err error
	var b []byte

	if err = r.ParseForm(); err != nil {
		logrus.WithError(err).Errorln("Failed to parse query")
		http.Error(w, "Failed to parse query", http.StatusBadRequest)
	}
	q, a, fields := r.Form.Get("q"), r.Form.Get("a"), r.Form["f"]

	// No error handling here. Using defaults if wrong params given
	offset, _ := strconv.Atoi(r.FormValue("offset"))
	maxResults, _ := strconv.Atoi(r.FormValue("maxResults"))
	if maxResults == 0 {
		maxResults = 10
	}

	if q == "" && a == "" {
		http.Error(w, "No search criteria provided", http.StatusBadRequest)
		return
	}

	var sr *dim.IndexResults
	l := logrus.WithFields(logrus.Fields{"query": q, "advanced_query": a, "fields": fields})
	l.Debugln("Searching image")
	if sr, err = i.SearchImages(q, a, fields, offset, maxResults); err != nil {
		http.Error(w, "An error occured while procesing your request", http.StatusInternalServerError)
		l.WithError(err).Errorln("Error occured when processing search")
		return
	}

	results := dim.SearchResults{NumResults: int(sr.Total), Query: q}
	l.WithField("#results", results.NumResults).Debugln("Found results")

	results.Results = buildResults(sr)

	if b, err = json.Marshal(results); err != nil {
		http.Error(w, "Failed to serialize the response", http.StatusInternalServerError)
		l.WithError(err).Errorln("Error occured while serializing search results")
	} else {
		fmt.Fprint(w, string(b))
	}
}

func buildResults(sr *dim.IndexResults) []dim.SearchResult {
	images := make([]dim.SearchResult, 0, sr.Total)
	for _, i := range sr.Images {
		images = append(images, imageToSearchResult(i))
	}
	return images
}

func imageToSearchResult(i *dim.IndexImage) dim.SearchResult {
	logrus.WithField("image", i).Debugln("Entering imageToSearchResult")
	result := dim.SearchResult{
		Name:         i.Name,
		Description:  i.Tag,
		Tag:          i.Tag,
		FullName:     i.FullName,
		Created:      i.Created,
		Label:        i.Label,
		Volumes:      i.Volumes,
		ExposedPorts: i.ExposedPorts,
		Env:          i.Env,
		Size:         i.Size,
	}

	return result
}

// DimHandlerFunc injects index into a HandlerFunc function
type DimHandlerFunc func(i dim.RegistryIndex, w http.ResponseWriter, r *http.Request)
<commit_msg>removing dead code<commit_after>// Copyright 2016
// 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 (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"time"

	"context"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/manifest/schema2"
	"github.com/docker/distribution/notifications"
	"github.com/mailgun/manners"
	"github.com/nhurel/dim/lib"
	"github.com/nhurel/dim/lib/environment"
)

// Server type handle  indexation of a docker registry and serves the search endpoint
type Server struct {
	*manners.GracefulServer
	index dim.RegistryIndex
}

// NewServer creates a new Server instance to listen on given port and use given index
func NewServer(port string, index dim.RegistryIndex, ctx context.Context) *Server {
	c := environment.Set(ctx, environment.StartTimeKey, time.Now())
	http.HandleFunc("/v1/search", handler(index, Search))
	http.HandleFunc("/dim/notify", handler(index, NotifyImageChange))
	http.HandleFunc("/dim/version", buildVersionHandler(c))
	http.HandleFunc("/v2/_catalog", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, "{}")
	})
	return &Server{manners.NewWithServer(&http.Server{Addr: port, Handler: http.DefaultServeMux}), index}
}

// Run starts the server instance
func (s *Server) Run() error {
	return s.ListenAndServe()
}

// Handler injects an index into an HandlerFunc
func handler(i dim.RegistryIndex, dhf DimHandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		dhf(i, w, r)
	}
}

func buildVersionHandler(ctx context.Context) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		Version(ctx, w, r)
	}
}

// Version return server info including info and uptime
func Version(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	info := dim.Info{}
	v := environment.Get(ctx, environment.VersionKey)
	if v != nil {
		info.Version = v.(string)
	}
	start := environment.Get(ctx, environment.StartTimeKey)
	if start != nil {
		info.Uptime = time.Now().Sub(start.(time.Time)).String()
	}
	if b, err := json.Marshal(info); err != nil {
		http.Error(w, "Failed to serialize the response", http.StatusInternalServerError)
		logrus.WithError(err).Errorln("Error occured while serializing server info")
	} else {
		fmt.Fprint(w, string(b))
	}
}

// NotifyImageChange handles docker registry events
func NotifyImageChange(i dim.RegistryIndex, w http.ResponseWriter, r *http.Request) {

	logrus.Debugln("Receiving event from registry")
	defer r.Body.Close()

	enveloppe := &notifications.Envelope{}

	if err := json.NewDecoder(r.Body).Decode(enveloppe); err != nil {
		logrus.WithError(err).Errorln("Failed to parse event")
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
	fmt.Fprintln(w, "Notiifcation handled !")

	logrus.WithField("enveloppe", enveloppe).Debugln("Processing event")

	for _, event := range enveloppe.Events {
		switch event.Action {
		case notifications.EventActionDelete:
			logrus.WithField("enveloppe", enveloppe).Infoln("Processing delete event")
			i.Submit(&dim.NotificationJob{Action: dim.DeleteAction, Digest: event.Target.Digest})
		case notifications.EventActionPush:
			if event.Target.MediaType == schema2.MediaTypeManifest {
				logrus.WithField("enveloppe", enveloppe).Infoln("Processing push event")
				i.Submit(&dim.NotificationJob{Action: dim.PushAction, Repository: event.Target.Repository, Tag: event.Target.Tag, Digest: event.Target.Digest})
			} else {
				logrus.WithField("mediatype", event.Target.MediaType).WithField("Event", event).Debugln("Event safely ignored because mediatype is unknown")
			}
		default:
			logrus.WithField("Action", event.Action).WithField("Event", event).Debugln("Event safely ignored")
		}
	}
}

// Search handles docker search request
func Search(i dim.RegistryIndex, w http.ResponseWriter, r *http.Request) {

	var err error
	var b []byte

	if err = r.ParseForm(); err != nil {
		logrus.WithError(err).Errorln("Failed to parse query")
		http.Error(w, "Failed to parse query", http.StatusBadRequest)
	}
	q, a, fields := r.Form.Get("q"), r.Form.Get("a"), r.Form["f"]

	// No error handling here. Using defaults if wrong params given
	offset, _ := strconv.Atoi(r.FormValue("offset"))
	maxResults, _ := strconv.Atoi(r.FormValue("maxResults"))
	if maxResults == 0 {
		maxResults = 10
	}

	if q == "" && a == "" {
		http.Error(w, "No search criteria provided", http.StatusBadRequest)
		return
	}

	var sr *dim.IndexResults
	l := logrus.WithFields(logrus.Fields{"query": q, "advanced_query": a, "fields": fields})
	l.Debugln("Searching image")
	if sr, err = i.SearchImages(q, a, fields, offset, maxResults); err != nil {
		http.Error(w, "An error occured while procesing your request", http.StatusInternalServerError)
		l.WithError(err).Errorln("Error occured when processing search")
		return
	}

	results := dim.SearchResults{NumResults: int(sr.Total), Query: q}
	l.WithField("#results", results.NumResults).Debugln("Found results")

	results.Results = buildResults(sr)

	if b, err = json.Marshal(results); err != nil {
		http.Error(w, "Failed to serialize the response", http.StatusInternalServerError)
		l.WithError(err).Errorln("Error occured while serializing search results")
	} else {
		fmt.Fprint(w, string(b))
	}
}

func buildResults(sr *dim.IndexResults) []dim.SearchResult {
	images := make([]dim.SearchResult, 0, sr.Total)
	for _, i := range sr.Images {
		images = append(images, imageToSearchResult(i))
	}
	return images
}

func imageToSearchResult(i *dim.IndexImage) dim.SearchResult {
	logrus.WithField("image", i).Debugln("Entering imageToSearchResult")
	result := dim.SearchResult{
		Name:         i.Name,
		Description:  i.Tag,
		Tag:          i.Tag,
		FullName:     i.FullName,
		Created:      i.Created,
		Label:        i.Label,
		Volumes:      i.Volumes,
		ExposedPorts: i.ExposedPorts,
		Env:          i.Env,
		Size:         i.Size,
	}

	return result
}

// DimHandlerFunc injects index into a HandlerFunc function
type DimHandlerFunc func(i dim.RegistryIndex, w http.ResponseWriter, r *http.Request)
<|endoftext|>


<commit_before>package server

// transfer2go agent server implementation
// Copyright (c) 2017 - Valentin Kuznetsov <vkuznet@gmail.com>

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/vkuznet/transfer2go/core"
	"github.com/vkuznet/transfer2go/utils"

	// web profiler, see https://golang.org/pkg/net/http/pprof
	_ "net/http/pprof"
)

// Config type holds server configuration
type Config struct {
	Name      string `json:"name"`      // agent name, aka site name
	Url       string `json:"url"`       // agent url
	Catalog   string `json:"catalog"`   // catalog file name, e.g. catalog.db
	Protocol  string `json:"protocol"`  // backend protocol, e.g. srmv2
	Backend   string `json:"backend"`   // backend, e.g. srm
	Tool      string `json:"tool"`      // backend tool, e.g. srmcp
	ToolOpts  string `json:"toolopts"`  // options for backend tool
	Mfile     string `json:"mfile"`     // metrics file name
	Minterval int64  `json:"minterval"` // metrics interval
	Staticdir string `json:"staticdir"` // static dir defines location of static files, e.g. sql,js templates
	Workers   int    `json:"workers"`   // number of workers
	QueueSize int    `json:"queuesize"` // total size of the queue
	Port      int    `json:"port"`      // port number given server runs on, default 8989
	Base      string `json:"base"`      // URL base path for agent server, it will be extracted from Url
	Register  string `json:"register"`  // remote agent URL to register
}

// String returns string representation of Config data type
func (c *Config) String() string {
	return fmt.Sprintf("<Config: name=%s url=%s port=%d base=%s catalog=%s protocol=%s backend=%s tool=%s opts=%s mfile=%s minterval=%d staticdir=%s workders=%d queuesize=%d register=%s>", c.Name, c.Url, c.Port, c.Base, c.Catalog, c.Protocol, c.Backend, c.Tool, c.ToolOpts, c.Mfile, c.Minterval, c.Staticdir, c.Workers, c.QueueSize, c.Register)
}

// AgentInfo data type
type AgentInfo struct {
	Agent string
	Alias string
}

// AgentProtocol data type
type AgentProtocol struct {
	Protocol string `json:"protocol"` // protocol name, e.g. srmv2
	Backend  string `json:"backend"`  // backend storage end-point, e.g. srm://cms-srm.cern.ch:8443/srm/managerv2?SFN=
	Tool     string `json:"tool"`     // actual executable, e.g. /usr/local/bin/srmcp
	ToolOpts string `json:"toolopts"` // options for backend tool
}

// globals used in server/handlers
var _myself, _alias, _protocol, _backend, _tool, _toolOpts string
var _agents map[string]string
var _config Config

// init
func init() {
	_agents = make(map[string]string)
}

// register a new (alias, agent) pair in agent (register)
func register(register, alias, agent string) error {
	log.Printf("Register %s as %s on %s\n", agent, alias, register)
	// register myself with another agent
	params := AgentInfo{Agent: _myself, Alias: _alias}
	data, err := json.Marshal(params)
	if err != nil {
		log.Println("ERROR, unable to marshal params", params)
	}
	url := fmt.Sprintf("%s/register", register)
	resp := utils.FetchResponse(url, data) // POST request
	// check return status code
	if resp.StatusCode != 200 {
		return fmt.Errorf("Response %s, error=%s", resp.Status, string(resp.Data))
	}
	return resp.Error
}

// helper function to register agent with all distributed agents
func registerAtAgents(aName string) {
	// register itself
	if _, ok := _agents[_alias]; ok {
		log.Fatal("ERROR unable to register", _alias, "at", _agents, "since this name already exists")
	}
	_agents[_alias] = _myself

	// now ask remote server for its list of agents and update internal map
	if aName != "" && len(aName) > 0 {
		err := register(aName, _alias, _myself) // submit remote registration of given agent name
		if err != nil {
			log.Fatal("ERROR Unable to register", _alias, _myself, "at", aName, err)
		}
		aurl := fmt.Sprintf("%s/agents", aName)
		resp := utils.FetchResponse(aurl, []byte{})
		var remoteAgents map[string]string
		e := json.Unmarshal(resp.Data, &remoteAgents)
		if e == nil {
			for key, val := range remoteAgents {
				if _, ok := _agents[key]; !ok {
					_agents[key] = val // register remote agent/alias pair internally
				}
			}
		}
	}

	// complete registration with other agents
	for alias, agent := range _agents {
		if agent == aName || alias == _alias {
			continue
		}
		register(agent, _alias, _myself) // submit remote registration of given agent name
	}

}

// Server implementation
func Server(config Config) {
	_config = config
	_myself = config.Url
	_alias = config.Name
	_protocol = config.Protocol
	_backend = config.Backend
	_tool = config.Tool
	_toolOpts = config.ToolOpts
	utils.STATICDIR = config.Staticdir
	arr := strings.Split(_myself, "/")
	base := ""
	if len(arr) > 3 {
		base = fmt.Sprintf("/%s", strings.Join(arr[3:], "/"))
	}
	port := "8989" // default port, the port here is a string type since we'll use it later in http.ListenAndServe
	if config.Port != 0 {
		port = fmt.Sprintf("%d", config.Port)
	}
	config.Base = base
	log.Println("Agent", config.String())

	// register self agent URI in remote agent and vice versa
	registerAtAgents(config.Register)

	// define catalog
	c, e := ioutil.ReadFile(config.Catalog)
	if e != nil {
		log.Fatalf("Unable to read catalog file, error=%v\n", e)
	}
	err := json.Unmarshal([]byte(c), &core.TFC)
	if err != nil {
		log.Fatalf("Unable to parse catalog JSON file, error=%v\n", err)
	}
	// open up Catalog DB
	dbtype := core.TFC.Type
	dburi := core.TFC.Uri // TODO: may be I need to change this based on DB Login/Password, check MySQL
	dbowner := core.TFC.Owner
	db, dberr := sql.Open(dbtype, dburi)
	defer db.Close()
	if dberr != nil {
		log.Fatalf("ERROR sql.Open, %v\n", dberr)
	}
	dberr = db.Ping()
	if dberr != nil {
		log.Fatalf("ERROR db.Ping, %v\n", dberr)
	}

	core.DB = db
	core.DBTYPE = dbtype
	core.DBSQL = core.LoadSQL(dbtype, dbowner)
	log.Println("Catalog", core.TFC)

	// define handlers
	http.HandleFunc(fmt.Sprintf("%s/status", base), StatusHandler)             // GET method
	http.HandleFunc(fmt.Sprintf("%s/agents", base), AgentsHandler)             // GET method
	http.HandleFunc(fmt.Sprintf("%s/files", base), FilesHandler)               // GET method
	http.HandleFunc(fmt.Sprintf("%s/reset", base), ResetHandler)               // GET method
	http.HandleFunc(fmt.Sprintf("%s/tfc", base), TFCHandler)                   // GET/POST method
	http.HandleFunc(fmt.Sprintf("%s/upload", base), UploadDataHandler)         // POST method
	http.HandleFunc(fmt.Sprintf("%s/request", base), RequestHandler)           // POST method
	http.HandleFunc(fmt.Sprintf("%s/register", base), RegisterAgentHandler)    // POST method
	http.HandleFunc(fmt.Sprintf("%s/protocol", base), RegisterProtocolHandler) // POST method
	http.HandleFunc(fmt.Sprintf("%s/", base), DefaultHandler)                  // GET method

	// initialize task dispatcher
	dispatcher := core.NewDispatcher(config.Workers, config.QueueSize, config.Mfile, config.Minterval)
	dispatcher.Run()
	log.Println("Start dispatcher with", config.Workers, "workers, queue size", config.QueueSize)

	// start server
	err = http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}
<commit_msg>Convert server to secure one; add AuthHandler<commit_after>package server

// transfer2go agent server implementation
// Copyright (c) 2017 - Valentin Kuznetsov <vkuznet@gmail.com>

import (
	"crypto/tls"
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/vkuznet/cmsauth"
	"github.com/vkuznet/transfer2go/core"
	"github.com/vkuznet/transfer2go/utils"

	// web profiler, see https://golang.org/pkg/net/http/pprof
	_ "net/http/pprof"
)

// Config type holds server configuration
type Config struct {
	Name      string `json:"name"`      // agent name, aka site name
	Url       string `json:"url"`       // agent url
	Catalog   string `json:"catalog"`   // catalog file name, e.g. catalog.db
	Protocol  string `json:"protocol"`  // backend protocol, e.g. srmv2
	Backend   string `json:"backend"`   // backend, e.g. srm
	Tool      string `json:"tool"`      // backend tool, e.g. srmcp
	ToolOpts  string `json:"toolopts"`  // options for backend tool
	Mfile     string `json:"mfile"`     // metrics file name
	Minterval int64  `json:"minterval"` // metrics interval
	Staticdir string `json:"staticdir"` // static dir defines location of static files, e.g. sql,js templates
	Workers   int    `json:"workers"`   // number of workers
	QueueSize int    `json:"queuesize"` // total size of the queue
	Port      int    `json:"port"`      // port number given server runs on, default 8989
	Base      string `json:"base"`      // URL base path for agent server, it will be extracted from Url
	Register  string `json:"register"`  // remote agent URL to register
	ServerKey string `json:"serverkey"` // server key file
	ServerCrt string `json:"servercrt"` // server crt file
}

// String returns string representation of Config data type
func (c *Config) String() string {
	return fmt.Sprintf("<Config: name=%s url=%s port=%d base=%s catalog=%s protocol=%s backend=%s tool=%s opts=%s mfile=%s minterval=%d staticdir=%s workders=%d queuesize=%d register=%s>", c.Name, c.Url, c.Port, c.Base, c.Catalog, c.Protocol, c.Backend, c.Tool, c.ToolOpts, c.Mfile, c.Minterval, c.Staticdir, c.Workers, c.QueueSize, c.Register)
}

// AgentInfo data type
type AgentInfo struct {
	Agent string
	Alias string
}

// AgentProtocol data type
type AgentProtocol struct {
	Protocol string `json:"protocol"` // protocol name, e.g. srmv2
	Backend  string `json:"backend"`  // backend storage end-point, e.g. srm://cms-srm.cern.ch:8443/srm/managerv2?SFN=
	Tool     string `json:"tool"`     // actual executable, e.g. /usr/local/bin/srmcp
	ToolOpts string `json:"toolopts"` // options for backend tool
}

// globals used in server/handlers
var _myself, _alias, _protocol, _backend, _tool, _toolOpts string
var _agents map[string]string
var _config Config
var _cmsAuth cmsauth.CMSAuth

// init
func init() {
	_agents = make(map[string]string)
}

// register a new (alias, agent) pair in agent (register)
func register(register, alias, agent string) error {
	log.Printf("Register %s as %s on %s\n", agent, alias, register)
	// register myself with another agent
	params := AgentInfo{Agent: _myself, Alias: _alias}
	data, err := json.Marshal(params)
	if err != nil {
		log.Println("ERROR, unable to marshal params", params)
	}
	url := fmt.Sprintf("%s/register", register)
	resp := utils.FetchResponse(url, data) // POST request
	// check return status code
	if resp.StatusCode != 200 {
		return fmt.Errorf("Response %s, error=%s", resp.Status, string(resp.Data))
	}
	return resp.Error
}

// helper function to register agent with all distributed agents
func registerAtAgents(aName string) {
	// register itself
	if _, ok := _agents[_alias]; ok {
		log.Fatal("ERROR unable to register", _alias, "at", _agents, "since this name already exists")
	}
	_agents[_alias] = _myself

	// now ask remote server for its list of agents and update internal map
	if aName != "" && len(aName) > 0 {
		err := register(aName, _alias, _myself) // submit remote registration of given agent name
		if err != nil {
			log.Fatal("ERROR Unable to register", _alias, _myself, "at", aName, err)
		}
		aurl := fmt.Sprintf("%s/agents", aName)
		resp := utils.FetchResponse(aurl, []byte{})
		var remoteAgents map[string]string
		e := json.Unmarshal(resp.Data, &remoteAgents)
		if e == nil {
			for key, val := range remoteAgents {
				if _, ok := _agents[key]; !ok {
					_agents[key] = val // register remote agent/alias pair internally
				}
			}
		}
	}

	// complete registration with other agents
	for alias, agent := range _agents {
		if agent == aName || alias == _alias {
			continue
		}
		register(agent, _alias, _myself) // submit remote registration of given agent name
	}

}

// Server implementation
func Server(config Config) {
	_config = config
	_myself = config.Url
	_alias = config.Name
	_protocol = config.Protocol
	_backend = config.Backend
	_tool = config.Tool
	_toolOpts = config.ToolOpts
	utils.STATICDIR = config.Staticdir
	arr := strings.Split(_myself, "/")
	base := ""
	if len(arr) > 3 {
		base = fmt.Sprintf("/%s", strings.Join(arr[3:], "/"))
	}
	port := "8989" // default port, the port here is a string type since we'll use it later in http.ListenAndServe
	if config.Port != 0 {
		port = fmt.Sprintf("%d", config.Port)
	}
	config.Base = base
	log.Println("Agent", config.String())

	// register self agent URI in remote agent and vice versa
	registerAtAgents(config.Register)

	// define catalog
	c, e := ioutil.ReadFile(config.Catalog)
	if e != nil {
		log.Fatalf("Unable to read catalog file, error=%v\n", e)
	}
	err := json.Unmarshal([]byte(c), &core.TFC)
	if err != nil {
		log.Fatalf("Unable to parse catalog JSON file, error=%v\n", err)
	}
	// open up Catalog DB
	dbtype := core.TFC.Type
	dburi := core.TFC.Uri // TODO: may be I need to change this based on DB Login/Password, check MySQL
	dbowner := core.TFC.Owner
	db, dberr := sql.Open(dbtype, dburi)
	defer db.Close()
	if dberr != nil {
		log.Fatalf("ERROR sql.Open, %v\n", dberr)
	}
	dberr = db.Ping()
	if dberr != nil {
		log.Fatalf("ERROR db.Ping, %v\n", dberr)
	}

	core.DB = db
	core.DBTYPE = dbtype
	core.DBSQL = core.LoadSQL(dbtype, dbowner)
	log.Println("Catalog", core.TFC)

	// define handlers
	http.HandleFunc(fmt.Sprintf("%s/", base), AuthHandler)

	// initialize task dispatcher
	dispatcher := core.NewDispatcher(config.Workers, config.QueueSize, config.Mfile, config.Minterval)
	dispatcher.Run()
	log.Println("Start dispatcher with", config.Workers, "workers, queue size", config.QueueSize)

	// start HTTPS server which require user certificates
	server := &http.Server{
		Addr: ":" + port,
		TLSConfig: &tls.Config{
			ClientAuth: tls.RequestClientCert,
		},
	}
	err = server.ListenAndServeTLS(config.ServerCrt, config.ServerKey)
	//     err = http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

	"github.com/heroku/busl/Godeps/_workspace/src/github.com/braintree/manners"
	"github.com/heroku/busl/Godeps/_workspace/src/github.com/gorilla/mux"
	"github.com/heroku/busl/Godeps/_workspace/src/github.com/heroku/rollbar"
	"github.com/heroku/busl/broker"
	"github.com/heroku/busl/util"
)

var gracefulServer *manners.GracefulServer

func init() {
	gracefulServer = manners.NewServer()
	gracefulServer.InnerServer.ReadTimeout = *util.HttpReadTimeout
	gracefulServer.InnerServer.WriteTimeout = *util.HttpWriteTimeout
}

func mkstream(w http.ResponseWriter, _ *http.Request) {
	registrar := broker.NewRedisRegistrar()
	uuid, err := util.NewUUID()
	if err != nil {
		http.Error(w, "Unable to create stream. Please try again.", http.StatusServiceUnavailable)
		rollbar.Error(rollbar.ERR, fmt.Errorf("unable to create new uuid for stream: %#v", err))
		util.CountWithData("mkstream.create.fail", 1, "error=%s", err)
		return
	}

	if err := registrar.Register(uuid); err != nil {
		http.Error(w, "Unable to create stream. Please try again.", http.StatusServiceUnavailable)
		rollbar.Error(rollbar.ERR, fmt.Errorf("unable to register stream: %#v", err))
		util.CountWithData("mkstream.create.fail", 1, "error=%s", err)
		return
	}

	util.Count("mkstream.create.success")
	io.WriteString(w, string(uuid))
}

func put(w http.ResponseWriter, r *http.Request) {
	registrar := broker.NewRedisRegistrar()

	if err := registrar.Register(key(r)); err != nil {
		http.Error(w, "Unable to create stream. Please try again.", http.StatusServiceUnavailable)
		rollbar.Error(rollbar.ERR, fmt.Errorf("unable to register stream: %#v", err))
		util.CountWithData("put.create.fail", 1, "error=%s", err)
		return
	}
	util.Count("put.create.success")
	w.WriteHeader(http.StatusCreated)
}

func health(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "OK")
}

func pub(w http.ResponseWriter, r *http.Request) {
	if !util.StringSliceUtil(r.TransferEncoding).Contains("chunked") {
		http.Error(w, "A chunked Transfer-Encoding header is required.", http.StatusBadRequest)
		return
	}

	writer, err := broker.NewWriter(key(r))
	if err != nil {
		handleError(w, r, err)
		return
	}
	defer writer.Close()

	body := bufio.NewReader(r.Body)
	defer r.Body.Close()

	_, err = io.Copy(writer, body)

	if err == io.ErrUnexpectedEOF {
		util.CountWithData("server.pub.read.eoferror", 1, "msg=\"%v\"", err.Error())
		return
	}

	if err != nil {
		log.Printf("%#v", err)
		http.Error(w, "Unhandled error, please try again.", http.StatusInternalServerError)
		rollbar.Error(rollbar.ERR, fmt.Errorf("unhandled error: %#v", err))
		return
	}

	done := make(chan struct{})

	go func() {
		for {
			select {
			case <-done:
				// Asynchronously upload the output to our defined storage backend.
				go storeOutput(key(r), requestURI(r))
				return
			case <-time.After(*util.StorageInterval):
				// Asynchronously upload the output to our defined storage backend.
				go storeOutput(key(r), requestURI(r))
			}
		}

	}()

	close(done)
}

func sub(w http.ResponseWriter, r *http.Request) {
	if _, ok := w.(http.Flusher); !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}

	rd, err := newReader(w, r)
	if rd != nil {
		defer rd.Close()
	}
	if err != nil {
		handleError(w, r, err)
		return
	}
	io.Copy(newWriteFlusher(w), rd)
}

func app() http.Handler {
	r := mux.NewRouter()

	r.HandleFunc("/health", addDefaultHeaders(health))

	// Legacy endpoint for creating the uuid `key` for you.
	r.HandleFunc("/streams", auth(addDefaultHeaders(mkstream)))

	// New `key` design for allowing any kind of id to be decided
	// by the caller (in this case, it mirrors what we have in S3).
	r.HandleFunc("/streams/{key:.+}", addDefaultHeaders(sub)).Methods("GET")
	r.HandleFunc("/streams/{key:.+}", addDefaultHeaders(pub)).Methods("POST")
	r.HandleFunc("/streams/{key:.+}", auth(addDefaultHeaders(put))).Methods("PUT")

	return logRequest(enforceHTTPS(r.ServeHTTP))
}

func Start(port string, shutdown <-chan struct{}) {
	log.Printf("http.start.port=%s\n", port)
	http.Handle("/", app())
	go listenForShutdown(shutdown)

	if err := gracefulServer.ListenAndServe(":"+port, nil); err != nil {
		log.Fatalf("server.server error=%v", err)
	}
}

func listenForShutdown(shutdown <-chan struct{}) {
	log.Println("http.graceful.await")
	<-shutdown
	log.Println("http.graceful.shutdown")
	gracefulServer.InnerServer.SetKeepAlivesEnabled(false) // TODO: Remove after merge of https://github.com/braintree/manners/pull/22
	gracefulServer.Shutdown <- true
}
<commit_msg>Fix the sequencing of the operations via @tt<commit_after>package server

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

	"github.com/heroku/busl/Godeps/_workspace/src/github.com/braintree/manners"
	"github.com/heroku/busl/Godeps/_workspace/src/github.com/gorilla/mux"
	"github.com/heroku/busl/Godeps/_workspace/src/github.com/heroku/rollbar"
	"github.com/heroku/busl/broker"
	"github.com/heroku/busl/util"
)

var gracefulServer *manners.GracefulServer

func init() {
	gracefulServer = manners.NewServer()
	gracefulServer.InnerServer.ReadTimeout = *util.HttpReadTimeout
	gracefulServer.InnerServer.WriteTimeout = *util.HttpWriteTimeout
}

func mkstream(w http.ResponseWriter, _ *http.Request) {
	registrar := broker.NewRedisRegistrar()
	uuid, err := util.NewUUID()
	if err != nil {
		http.Error(w, "Unable to create stream. Please try again.", http.StatusServiceUnavailable)
		rollbar.Error(rollbar.ERR, fmt.Errorf("unable to create new uuid for stream: %#v", err))
		util.CountWithData("mkstream.create.fail", 1, "error=%s", err)
		return
	}

	if err := registrar.Register(uuid); err != nil {
		http.Error(w, "Unable to create stream. Please try again.", http.StatusServiceUnavailable)
		rollbar.Error(rollbar.ERR, fmt.Errorf("unable to register stream: %#v", err))
		util.CountWithData("mkstream.create.fail", 1, "error=%s", err)
		return
	}

	util.Count("mkstream.create.success")
	io.WriteString(w, string(uuid))
}

func put(w http.ResponseWriter, r *http.Request) {
	registrar := broker.NewRedisRegistrar()

	if err := registrar.Register(key(r)); err != nil {
		http.Error(w, "Unable to create stream. Please try again.", http.StatusServiceUnavailable)
		rollbar.Error(rollbar.ERR, fmt.Errorf("unable to register stream: %#v", err))
		util.CountWithData("put.create.fail", 1, "error=%s", err)
		return
	}
	util.Count("put.create.success")
	w.WriteHeader(http.StatusCreated)
}

func health(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "OK")
}

func pub(w http.ResponseWriter, r *http.Request) {
	if !util.StringSliceUtil(r.TransferEncoding).Contains("chunked") {
		http.Error(w, "A chunked Transfer-Encoding header is required.", http.StatusBadRequest)
		return
	}

	writer, err := broker.NewWriter(key(r))
	if err != nil {
		handleError(w, r, err)
		return
	}
	defer writer.Close()

	body := bufio.NewReader(r.Body)
	defer r.Body.Close()

	done := make(chan struct{})

	go func() {
		for {
			select {
			case <-done:
				// Asynchronously upload the output to our defined storage backend.
				go storeOutput(key(r), requestURI(r))
				return
			case <-time.After(*util.StorageInterval):
				// Asynchronously upload the output to our defined storage backend.
				go storeOutput(key(r), requestURI(r))
			}
		}

	}()

	_, err = io.Copy(writer, body)

	if err == io.ErrUnexpectedEOF {
		util.CountWithData("server.pub.read.eoferror", 1, "msg=\"%v\"", err.Error())
		return
	}

	if err != nil {
		log.Printf("%#v", err)
		http.Error(w, "Unhandled error, please try again.", http.StatusInternalServerError)
		rollbar.Error(rollbar.ERR, fmt.Errorf("unhandled error: %#v", err))
		return
	}

	close(done)
}

func sub(w http.ResponseWriter, r *http.Request) {
	if _, ok := w.(http.Flusher); !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}

	rd, err := newReader(w, r)
	if rd != nil {
		defer rd.Close()
	}
	if err != nil {
		handleError(w, r, err)
		return
	}
	io.Copy(newWriteFlusher(w), rd)
}

func app() http.Handler {
	r := mux.NewRouter()

	r.HandleFunc("/health", addDefaultHeaders(health))

	// Legacy endpoint for creating the uuid `key` for you.
	r.HandleFunc("/streams", auth(addDefaultHeaders(mkstream)))

	// New `key` design for allowing any kind of id to be decided
	// by the caller (in this case, it mirrors what we have in S3).
	r.HandleFunc("/streams/{key:.+}", addDefaultHeaders(sub)).Methods("GET")
	r.HandleFunc("/streams/{key:.+}", addDefaultHeaders(pub)).Methods("POST")
	r.HandleFunc("/streams/{key:.+}", auth(addDefaultHeaders(put))).Methods("PUT")

	return logRequest(enforceHTTPS(r.ServeHTTP))
}

func Start(port string, shutdown <-chan struct{}) {
	log.Printf("http.start.port=%s\n", port)
	http.Handle("/", app())
	go listenForShutdown(shutdown)

	if err := gracefulServer.ListenAndServe(":"+port, nil); err != nil {
		log.Fatalf("server.server error=%v", err)
	}
}

func listenForShutdown(shutdown <-chan struct{}) {
	log.Println("http.graceful.await")
	<-shutdown
	log.Println("http.graceful.shutdown")
	gracefulServer.InnerServer.SetKeepAlivesEnabled(false) // TODO: Remove after merge of https://github.com/braintree/manners/pull/22
	gracefulServer.Shutdown <- true
}
<|endoftext|>