<commit_before>package main

import (
	"bytes"
	"fmt"
	"image"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/PuerkitoBio/throttled"
	"github.com/PuerkitoBio/throttled/store"
	"github.com/codegangsta/cli"
	"github.com/go-martini/martini"
	"github.com/martini-contrib/binding"
)

// UploadForm is a form structure to use when an image is POSTed to the server
type UploadForm struct {
	PhotoUpload *multipart.FileHeader `form:"photo"`
}

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {
	// Set up logging for CLI
	log.SetPrefix("")
	log.SetFlags(0) // Remove the timestamp

	// Connect to redis
	err := redisInit()
	if err != nil {
		log.Println("Connecting to redis failed", err)
	}

	app := cli.NewApp()
	app.Name = "pixlserv"
	app.Usage = "transform and serve images"
	app.Version = "1.0"
	app.Commands = []cli.Command{
		{
			Name:  "run",
			Usage: "Runs the server (run [config-file])",
			Action: func(c *cli.Context) {
				// Set up logging for server
				log.SetPrefix("[pixlserv] ")

				if len(c.Args()) < 1 {
					log.Println("You need to provide a path to a config file")
					return
				}
				configFilePath := c.Args().First()

				// Initialise configuration
				err := configInit(configFilePath)
				if err != nil {
					log.Println("Configuration reading failed:", err)
					return
				}
				log.Printf("Running with config: %+v", Config)

				// Initialise authentication
				err = authInit()
				if err != nil {
					log.Println("Authentication initialisation failed:", err)
					return
				}

				// Initialise storage
				err = storageInit()
				if err != nil {
					log.Println("Storage initialisation failed:", err)
					return
				}

				// Run the server
				m := martini.Classic()
				if Config.throttlingRate > 0 {
					m.Use(throttler(Config.throttlingRate))
				}
				m.Get("/", func() string {
					return "It works!"
				})
				m.Get("/((?P<apikey>[A-Z0-9]+)/)?image/:parameters/**", transformationHandler)
				m.Post("/((?P<apikey>[A-Z0-9]+)/)?upload", binding.MultipartForm(UploadForm{}), uploadHandler)
				go m.Run()

				// Wait for when the program is terminated
				ch := make(chan os.Signal)
				signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
				<-ch

				// Clean up
				redisCleanUp()
				storageCleanUp()
			},
		},
		{
			Name:  "api-key",
			Usage: "Manages API keys",
			Subcommands: []cli.Command{
				{
					Name:  "add",
					Usage: "Adds a new one",
					Action: func(c *cli.Context) {
						key, err := generateKey()
						if err != nil {
							log.Println("Adding a new API key failed, please try again")
							return
						}

						log.Println("Key added:", key)
					},
				},
				{
					Name:  "info",
					Usage: "Shows information about a key (info [key])",
					Action: func(c *cli.Context) {
						if len(c.Args()) < 1 {
							log.Println("You need to provide an existing key")
							return
						}
						key := c.Args().First()
						permissions, err := infoAboutKey(key)
						if err != nil {
							log.Println(err.Error())
							return
						}
						log.Println("Key:", key)
						log.Println("Permissions:", permissions)
					},
				},
				{
					Name:  "list",
					Usage: "Shows all keys",
					Action: func(c *cli.Context) {
						keys, err := listKeys()
						if err != nil {
							log.Println("Retrieving the list of all keys failed")
							return
						}

						log.Println("Keys:", keys)
					},
				},
				{
					Name:  "modify",
					Usage: "Modifies permissions for a key (modify [key] [add/remove] [" + authPermissionsOptions() + "])",
					Action: func(c *cli.Context) {
						if len(c.Args()) < 3 {
							log.Println("You need to provide an existing key, operation and a permission")
							return
						}
						key := c.Args().First()
						err := modifyKey(key, c.Args()[1], c.Args()[2])
						if err != nil {
							log.Println(err.Error())
							return
						}
						log.Println("The key has been updated")
					},
				},
				{
					Name:  "remove",
					Usage: "Removes an existing key (remove [key])",
					Action: func(c *cli.Context) {
						if len(c.Args()) < 1 {
							log.Println("You need to provide an existing key")
							return
						}
						err := removeKey(c.Args().First())
						if err != nil {
							log.Println(err.Error())
							return
						}
						log.Println("The key was successfully removed")
					},
				},
			},
		},
	}

	app.Run(os.Args)
}

func transformationHandler(params martini.Params) (int, string) {
	if !hasPermission(params["apikey"], GetPermission) {
		return http.StatusUnauthorized, ""
	}

	var transformation Transformation
	transformationName := parseTransformationName(params["parameters"])
	if transformationName != "" {
		var ok bool
		transformation, ok = Config.transformations[transformationName]
		if !ok {
			return http.StatusBadRequest, "Unknown transformation: " + transformationName
		}
	} else if Config.allowCustomTransformations {
		parameters, err := parseParameters(params["parameters"])
		if err != nil {
			return http.StatusBadRequest, err.Error()
		}
		transformation = Transformation{&parameters, nil}
	} else {
		return http.StatusBadRequest, "Custom transformations not allowed"
	}
	baseImagePath, scale := parseBasePathAndScale(params["_1"])
	if Config.allowCustomScale {
		parameters := transformation.params.WithScale(scale)
		transformation.params = &parameters
	}

	// Check if the image with the given parameters already exists
	// and return it
	fullImagePath, _ := createFilePath(baseImagePath, transformation.params)
	img, format, err := loadFromCache(fullImagePath)
	if err == nil {
		var buffer bytes.Buffer
		writeImage(img, format, &buffer)

		return http.StatusOK, buffer.String()
	}

	// Load the original image and process it
	if !imageExists(baseImagePath) {
		return http.StatusNotFound, "Image not found: " + baseImagePath
	}

	img, format, err = loadImage(baseImagePath)
	if err != nil {
		return http.StatusInternalServerError, err.Error()
	}

	imgNew := transformCropAndResize(img, &transformation)

	var buffer bytes.Buffer
	err = writeImage(imgNew, format, &buffer)
	if err != nil {
		log.Println("Writing an image to the response failed:", err)
	}

	// Cache the image asynchronously to speed up the response
	go func() {
		err = addToCache(fullImagePath, imgNew, format)
		if err != nil {
			log.Println("Saving an image to cache failed:", err)
		}
	}()

	return http.StatusOK, buffer.String()
}

func uploadHandler(params martini.Params, uf UploadForm) (int, string) {
	if !hasPermission(params["apikey"], UploadPermission) {
		return http.StatusUnauthorized, ""
	}

	file, err := uf.PhotoUpload.Open()
	if err != nil {
		return http.StatusBadRequest, err.Error()
	}

	limit := io.LimitReader(file, int64(Config.uploadMaxFileSize+1))
	data, err := ioutil.ReadAll(limit)
	if err != nil {
		return http.StatusBadRequest, err.Error()
	}
	if len(data) > Config.uploadMaxFileSize {
		return http.StatusBadRequest, "max file size exceeded"
	}

	img, format, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return http.StatusBadRequest, err.Error()
	}

	defer file.Close()

	// Not a big fan of .jpeg file extensions
	randomInt := rand.Intn(1000)
	baseImagePath := fmt.Sprintf("%d-%d.%s", time.Now().Unix(), randomInt, strings.Replace(format, "jpeg", "jpg", 1))
	log.Printf("Uploading %s", baseImagePath)

	// Eager transformations
	eagerlyTransform := func() {
		if len(Config.eagerTransformations) > 0 {
			for _, transformation := range Config.eagerTransformations {
				imgNew := transformCropAndResize(img, &transformation)
				fullImagePath, _ := createFilePath(baseImagePath, transformation.params)
				addToCache(fullImagePath, imgNew, format)
			}
		}
	}

	if Config.asyncUploads {
		go func() {
			_, err := saveImage(img, format, baseImagePath)
			if err != nil {
				log.Println("Error saving image:", err)
				return
			}
			go eagerlyTransform()
		}()
	} else {
		_, err := saveImage(img, format, baseImagePath)
		if err != nil {
			return http.StatusInternalServerError, "Error saving image: " + err.Error()
		}
		go eagerlyTransform()
	}

	return http.StatusOK, ""
}

func throttler(perMinRate int) http.Handler {
	t := throttled.RateLimit(throttled.PerMin(perMinRate), &throttled.VaryBy{RemoteAddr: true}, store.NewMemStore(1000))
	return t.Throttle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Nothing needed here
	}))
}
<commit_msg>Upload response returns file name of the uploaded file<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"image"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"

	"github.com/PuerkitoBio/throttled"
	"github.com/PuerkitoBio/throttled/store"
	"github.com/codegangsta/cli"
	"github.com/go-martini/martini"
	"github.com/martini-contrib/binding"
)

// UploadForm is a form structure to use when an image is POSTed to the server
type UploadForm struct {
	PhotoUpload *multipart.FileHeader `form:"photo"`
}

var (
	uploadUrlRe = regexp.MustCompile("/upload$")
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {
	// Set up logging for CLI
	log.SetPrefix("")
	log.SetFlags(0) // Remove the timestamp

	// Connect to redis
	err := redisInit()
	if err != nil {
		log.Println("Connecting to redis failed", err)
	}

	app := cli.NewApp()
	app.Name = "pixlserv"
	app.Usage = "transform and serve images"
	app.Version = "1.0"
	app.Commands = []cli.Command{
		{
			Name:  "run",
			Usage: "Runs the server (run [config-file])",
			Action: func(c *cli.Context) {
				// Set up logging for server
				log.SetPrefix("[pixlserv] ")

				if len(c.Args()) < 1 {
					log.Println("You need to provide a path to a config file")
					return
				}
				configFilePath := c.Args().First()

				// Initialise configuration
				err := configInit(configFilePath)
				if err != nil {
					log.Println("Configuration reading failed:", err)
					return
				}
				log.Printf("Running with config: %+v", Config)

				// Initialise authentication
				err = authInit()
				if err != nil {
					log.Println("Authentication initialisation failed:", err)
					return
				}

				// Initialise storage
				err = storageInit()
				if err != nil {
					log.Println("Storage initialisation failed:", err)
					return
				}

				// Run the server
				m := martini.Classic()
				if Config.throttlingRate > 0 {
					m.Use(throttler(Config.throttlingRate))
				}
				m.Use(func(res http.ResponseWriter, req *http.Request) {
					if uploadUrlRe.MatchString(req.URL.Path) {
						// The upload handler returns JSON
						res.Header().Set("Content-Type", "application/json")
					}
				})
				m.Get("/", func() string {
					return "It works!"
				})
				m.Get("/((?P<apikey>[A-Z0-9]+)/)?image/:parameters/**", transformationHandler)
				m.Post("/((?P<apikey>[A-Z0-9]+)/)?upload", binding.MultipartForm(UploadForm{}), uploadHandler)
				go m.Run()

				// Wait for when the program is terminated
				ch := make(chan os.Signal)
				signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
				<-ch

				// Clean up
				redisCleanUp()
				storageCleanUp()
			},
		},
		{
			Name:  "api-key",
			Usage: "Manages API keys",
			Subcommands: []cli.Command{
				{
					Name:  "add",
					Usage: "Adds a new one",
					Action: func(c *cli.Context) {
						key, err := generateKey()
						if err != nil {
							log.Println("Adding a new API key failed, please try again")
							return
						}

						log.Println("Key added:", key)
					},
				},
				{
					Name:  "info",
					Usage: "Shows information about a key (info [key])",
					Action: func(c *cli.Context) {
						if len(c.Args()) < 1 {
							log.Println("You need to provide an existing key")
							return
						}
						key := c.Args().First()
						permissions, err := infoAboutKey(key)
						if err != nil {
							log.Println(err.Error())
							return
						}
						log.Println("Key:", key)
						log.Println("Permissions:", permissions)
					},
				},
				{
					Name:  "list",
					Usage: "Shows all keys",
					Action: func(c *cli.Context) {
						keys, err := listKeys()
						if err != nil {
							log.Println("Retrieving the list of all keys failed")
							return
						}

						log.Println("Keys:", keys)
					},
				},
				{
					Name:  "modify",
					Usage: "Modifies permissions for a key (modify [key] [add/remove] [" + authPermissionsOptions() + "])",
					Action: func(c *cli.Context) {
						if len(c.Args()) < 3 {
							log.Println("You need to provide an existing key, operation and a permission")
							return
						}
						key := c.Args().First()
						err := modifyKey(key, c.Args()[1], c.Args()[2])
						if err != nil {
							log.Println(err.Error())
							return
						}
						log.Println("The key has been updated")
					},
				},
				{
					Name:  "remove",
					Usage: "Removes an existing key (remove [key])",
					Action: func(c *cli.Context) {
						if len(c.Args()) < 1 {
							log.Println("You need to provide an existing key")
							return
						}
						err := removeKey(c.Args().First())
						if err != nil {
							log.Println(err.Error())
							return
						}
						log.Println("The key was successfully removed")
					},
				},
			},
		},
	}

	app.Run(os.Args)
}

func transformationHandler(params martini.Params) (int, string) {
	if !hasPermission(params["apikey"], GetPermission) {
		return http.StatusUnauthorized, ""
	}

	var transformation Transformation
	transformationName := parseTransformationName(params["parameters"])
	if transformationName != "" {
		var ok bool
		transformation, ok = Config.transformations[transformationName]
		if !ok {
			return http.StatusBadRequest, "Unknown transformation: " + transformationName
		}
	} else if Config.allowCustomTransformations {
		parameters, err := parseParameters(params["parameters"])
		if err != nil {
			return http.StatusBadRequest, err.Error()
		}
		transformation = Transformation{&parameters, nil}
	} else {
		return http.StatusBadRequest, "Custom transformations not allowed"
	}
	baseImagePath, scale := parseBasePathAndScale(params["_1"])
	if Config.allowCustomScale {
		parameters := transformation.params.WithScale(scale)
		transformation.params = &parameters
	}

	// Check if the image with the given parameters already exists
	// and return it
	fullImagePath, _ := createFilePath(baseImagePath, transformation.params)
	img, format, err := loadFromCache(fullImagePath)
	if err == nil {
		var buffer bytes.Buffer
		writeImage(img, format, &buffer)

		return http.StatusOK, buffer.String()
	}

	// Load the original image and process it
	if !imageExists(baseImagePath) {
		return http.StatusNotFound, "Image not found: " + baseImagePath
	}

	img, format, err = loadImage(baseImagePath)
	if err != nil {
		return http.StatusInternalServerError, err.Error()
	}

	imgNew := transformCropAndResize(img, &transformation)

	var buffer bytes.Buffer
	err = writeImage(imgNew, format, &buffer)
	if err != nil {
		log.Println("Writing an image to the response failed:", err)
	}

	// Cache the image asynchronously to speed up the response
	go func() {
		err = addToCache(fullImagePath, imgNew, format)
		if err != nil {
			log.Println("Saving an image to cache failed:", err)
		}
	}()

	return http.StatusOK, buffer.String()
}

type UploadResponse struct {
	Status       string `json:"status"`
	ErrorMessage string `json:"errorMessage"`
	ImagePath    string `json:"imagePath"`
}

func uploadResponse(response UploadResponse) string {
	str, err := json.Marshal(response)
	if err != nil {
		log.Printf("Error constructing JSON response for %v", response)
		return "{\"status\": \"error\", \"errorMessage\": \"server error\"}"
	}
	return string(str[:])
}

func uploadError(errorMessage string) string {
	return uploadResponse(UploadResponse{"error", errorMessage, ""})
}

func uploadSuccess(imagePath string) string {
	return uploadResponse(UploadResponse{"ok", "", imagePath})
}

func uploadHandler(params martini.Params, uf UploadForm) (int, string) {
	if !hasPermission(params["apikey"], UploadPermission) {
		return http.StatusUnauthorized, uploadError("API key invalid or missing")
	}

	file, err := uf.PhotoUpload.Open()
	if err != nil {
		return http.StatusBadRequest, uploadError(err.Error())
	}

	limit := io.LimitReader(file, int64(Config.uploadMaxFileSize+1))
	data, err := ioutil.ReadAll(limit)
	if err != nil {
		return http.StatusBadRequest, uploadError(err.Error())
	}
	if len(data) > Config.uploadMaxFileSize {
		return http.StatusBadRequest, uploadError("max file size exceeded")
	}

	img, format, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return http.StatusBadRequest, uploadError(err.Error())
	}

	defer file.Close()

	// Not a big fan of .jpeg file extensions
	randomInt := rand.Intn(1000)
	baseImagePath := fmt.Sprintf("%d-%d.%s", time.Now().Unix(), randomInt, strings.Replace(format, "jpeg", "jpg", 1))
	log.Printf("Uploading %s", baseImagePath)

	// Eager transformations
	eagerlyTransform := func() {
		if len(Config.eagerTransformations) > 0 {
			for _, transformation := range Config.eagerTransformations {
				imgNew := transformCropAndResize(img, &transformation)
				fullImagePath, _ := createFilePath(baseImagePath, transformation.params)
				addToCache(fullImagePath, imgNew, format)
			}
		}
	}

	if Config.asyncUploads {
		go func() {
			_, err := saveImage(img, format, baseImagePath)
			if err != nil {
				log.Println("Error saving image:", err)
				return
			}
			go eagerlyTransform()
		}()
	} else {
		_, err := saveImage(img, format, baseImagePath)
		if err != nil {
			return http.StatusInternalServerError, uploadError("error saving image: " + err.Error())
		}
		go eagerlyTransform()
	}

	return http.StatusOK, uploadSuccess(baseImagePath)
}

func throttler(perMinRate int) http.Handler {
	t := throttled.RateLimit(throttled.PerMin(perMinRate), &throttled.VaryBy{RemoteAddr: true}, store.NewMemStore(1000))
	return t.Throttle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Nothing needed here
	}))
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"time"

	"github.com/streadway/amqp"
)

// a generous 60 seconds to apply raft commands
const raftMaxTime = time.Duration(60) * time.Second

// Server is the delayd server. It handles the server lifecycle (startup, clean shutdown)
type Server struct {
	sender   *AmqpSender
	receiver *AmqpReceiver
	storage  *Storage
	raft     *Raft
	timer    *Timer
}

// Run initializes the Server from a Config, and begins its main loop.
func (s *Server) Run(c Config) {
	log.Println("Starting delayd")

	log.Println("Creating data dir: ", c.DataDir)
	err := os.MkdirAll(c.DataDir, 0755)

	if err != nil {
		log.Fatal("Error creating data dir: ", err)
	}

	s.receiver, err = NewAmqpReceiver(c.Amqp)
	if err != nil {
		log.Fatal("Could not initialize receiver: ", err)
	}

	s.sender, err = NewAmqpSender(c.Amqp.URL)
	if err != nil {
		log.Fatal("Could not initialize sender: ", err)
	}

	s.storage, err = NewStorage(c.DataDir)
	if err != nil {
		log.Fatal("Could not initialize storage backend: ", err)
	}

	s.raft, err = NewRaft(c.DataDir, s.storage)
	if err != nil {
		log.Fatal("Could not initialize raft: ", err)
	}

	s.timer = NewTimer(s.timerSend)
	go s.observeLeaderChanges()

	for {
		eWrapper, ok := <-s.receiver.C
		entry := eWrapper.Entry
		// XXX cleanup needed here before exit
		if !ok {
			log.Fatal("Receiver Consumption failed!")
		}

		log.Println("Got entry: ", entry)
		b, err := entry.ToBytes()
		if err != nil {
			log.Println("Error encoding entry", err)
			continue
		}

		err = s.raft.Add(b, raftMaxTime)
		if err != nil {
			eWrapper.Done(false)
		}
		s.timer.Reset(entry.SendAt, false)

		eWrapper.Done(true)
	}
}

// Stop shuts down the Server cleanly. Order of the Close calls is important.
func (s *Server) Stop() {
	log.Println("Shutting down gracefully.")

	s.receiver.Close()
	s.storage.Close()
	s.sender.Close()

	s.raft.Close()

	log.Println("Terminated.")
}

func (s *Server) resetTimer() {
	ok, t, err := s.storage.NextTime()
	if err != nil {
		log.Fatal("Could not read initial send time from storage: ", err)
	}
	if ok {
		s.timer.Reset(t, true)
	}
}

// Listen for changes to raft for when we transition to and from the leader state
// and react accordingly.
func (s *Server) observeLeaderChanges() {
	for isLeader := range s.raft.LeaderCh() {
		switch isLeader {
		case true:
			log.Println("Became raft leader")
			s.resetTimer()
			err := s.receiver.Start()
			if err != nil {
				log.Fatal("Error while starting receiver: ", err)
			}

		case false:
			log.Println("Lost raft leadership")
			s.timer.Pause()
			err := s.receiver.Pause()
			if err != nil {
				log.Fatal("Error while starting receiver: ", err)
			}
		}
	}
}

func (s *Server) timerSend(t time.Time) (next time.Time, ok bool) {
	uuids, entries, err := s.storage.Get(t)
	if err != nil {
		log.Fatal("Could not read entries from db: ", err)
	}

	log.Printf("Sending %d entries\n", len(entries))
	for i, e := range entries {
		err = s.sender.Send(e)

		// error 504 code means that the exchange we were trying
		// to send on didnt exist.  In the case of delayd this usually
		// means that a consumer didn't set up the exchange they wish
		// to be notified on.  We do not attempt to make this for them,
		// as we don't know what exchange options they would want, we
		// simply drop this message, other errors are fatal

		if err, ok := err.(*amqp.Error); ok {
			if err.Code != 504 {
				log.Fatal("Could not send entry: ", err)
			} else {
				log.Printf("channel/connection not set up for exchange `%s`, message will be deleted", e.Target)
			}
		}

		err = s.raft.Remove(uuids[i], raftMaxTime)
		if err != nil {
			// XXX abort comitting for raft here, so it can retry.
			log.Fatal("Could not remove entry from db: ", err)
		}
	}

	// ensure everyone is up to date
	err = s.raft.SyncAll()
	if err != nil {
		log.Fatal("Could not sync all cluster nodes: ", err)
	}

	ok, next, err = s.storage.NextTime()
	if err != nil {
		log.Fatal("Could not read next time from db: ", err)
	}

	return
}
<commit_msg>Raft leadership loss is not fatal when emitting<commit_after>package main

import (
	"log"
	"os"
	"time"

	"github.com/streadway/amqp"
)

// a generous 60 seconds to apply raft commands
const raftMaxTime = time.Duration(60) * time.Second

// Server is the delayd server. It handles the server lifecycle (startup, clean shutdown)
type Server struct {
	sender   *AmqpSender
	receiver *AmqpReceiver
	storage  *Storage
	raft     *Raft
	timer    *Timer
}

// Run initializes the Server from a Config, and begins its main loop.
func (s *Server) Run(c Config) {
	log.Println("Starting delayd")

	log.Println("Creating data dir: ", c.DataDir)
	err := os.MkdirAll(c.DataDir, 0755)

	if err != nil {
		log.Fatal("Error creating data dir: ", err)
	}

	s.receiver, err = NewAmqpReceiver(c.Amqp)
	if err != nil {
		log.Fatal("Could not initialize receiver: ", err)
	}

	s.sender, err = NewAmqpSender(c.Amqp.URL)
	if err != nil {
		log.Fatal("Could not initialize sender: ", err)
	}

	s.storage, err = NewStorage(c.DataDir)
	if err != nil {
		log.Fatal("Could not initialize storage backend: ", err)
	}

	s.raft, err = NewRaft(c.DataDir, s.storage)
	if err != nil {
		log.Fatal("Could not initialize raft: ", err)
	}

	s.timer = NewTimer(s.timerSend)
	go s.observeLeaderChanges()

	for {
		eWrapper, ok := <-s.receiver.C
		entry := eWrapper.Entry
		// XXX cleanup needed here before exit
		if !ok {
			log.Fatal("Receiver Consumption failed!")
		}

		log.Println("Got entry: ", entry)
		b, err := entry.ToBytes()
		if err != nil {
			log.Println("Error encoding entry", err)
			continue
		}

		err = s.raft.Add(b, raftMaxTime)
		if err != nil {
			eWrapper.Done(false)
		}
		s.timer.Reset(entry.SendAt, false)

		eWrapper.Done(true)
	}
}

// Stop shuts down the Server cleanly. Order of the Close calls is important.
func (s *Server) Stop() {
	log.Println("Shutting down gracefully.")

	s.receiver.Close()
	s.storage.Close()
	s.sender.Close()

	s.raft.Close()

	log.Println("Terminated.")
}

func (s *Server) resetTimer() {
	ok, t, err := s.storage.NextTime()
	if err != nil {
		log.Fatal("Could not read initial send time from storage: ", err)
	}
	if ok {
		s.timer.Reset(t, true)
	}
}

// Listen for changes to raft for when we transition to and from the leader state
// and react accordingly.
func (s *Server) observeLeaderChanges() {
	for isLeader := range s.raft.LeaderCh() {
		switch isLeader {
		case true:
			log.Println("Became raft leader")
			s.resetTimer()
			err := s.receiver.Start()
			if err != nil {
				log.Fatal("Error while starting receiver: ", err)
			}

		case false:
			log.Println("Lost raft leadership")
			s.timer.Pause()
			err := s.receiver.Pause()
			if err != nil {
				log.Fatal("Error while starting receiver: ", err)
			}
		}
	}
}

func (s *Server) timerSend(t time.Time) (next time.Time, ok bool) {
	uuids, entries, err := s.storage.Get(t)
	if err != nil {
		log.Fatal("Could not read entries from db: ", err)
	}

	log.Printf("Sending %d entries\n", len(entries))
	for i, e := range entries {
		err = s.sender.Send(e)

		// error 504 code means that the exchange we were trying
		// to send on didnt exist.  In the case of delayd this usually
		// means that a consumer didn't set up the exchange they wish
		// to be notified on.  We do not attempt to make this for them,
		// as we don't know what exchange options they would want, we
		// simply drop this message, other errors are fatal

		if err, ok := err.(*amqp.Error); ok {
			if err.Code != 504 {
				log.Fatal("Could not send entry: ", err)
			} else {
				log.Printf("channel/connection not set up for exchange `%s`, message will be deleted", e.Target)
			}
		}

		err = s.raft.Remove(uuids[i], raftMaxTime)
		if err != nil {
			// This node is no longer the leader. give up on other amqp sends,
			// and scheduling the next emission
			log.Printf("Lost raft leadership during remove. AMQP send will be a duplicate. uuid=%x\n", uuids[i])
			break
		}
	}

	if err != nil {
		// don't reschedule
		ok = false
		return
	}

	// ensure everyone is up to date
	err = s.raft.SyncAll()
	if err != nil {
		log.Printf("Lost raft leadership during sync after send.")
		ok = false
		return
	}

	ok, next, err = s.storage.NextTime()
	if err != nil {
		log.Fatal("Could not read next time from db: ", err)
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
)

var port = flag.Int("port", 8080, "which port to listen on")
var authApiKey string
var dbUrl string

var rootMetrics = map[string]string{
	"system.load.1":       "system.load.1",
	"system.load.5":       "system.load.5",
	"system.load.15":      "system.load.15",
	"system.load.norm.1":  "system.load.norm.1",
	"system.load.norm.5":  "system.load.norm.5",
	"system.load.norm.15": "system.load.norm.15",

	"cpuIdle":   "system.cpu.idle",
	"cpuUser":   "system.cpu.user",
	"cpuWait":   "system.cpu.iowait",
	"cpuSystem": "system.cpu.system",
	"cpuStolen": "system.cpu.stolen",

	"memBuffers":       "system.mem.buffered",
	"memPhysPctUsable": "system.mem.pct_usable",
	"memShared":        "system.mem.shared",
	"memPhysTotal":     "system.mem.total",
	"memPhysUsable":    "system.mem.usable",
	"memCached":        "system.mem.cached",
	"memPhysUsed":      "system.mem.used",
	"memPhysFree":      "system.mem.free",

	"memSwapFree":    "system.swap.free",
	"memSwapTotal":   "system.swap.total",
	"memSwapUsed":    "system.swap.used",
	"memSwapPctFree": "system.swap.pct_free",
}

type Metric struct {
	Name    string          `json:"name"`
	Columns []string        `json:"columns"`
	Points  [][]interface{} `json:"points"`
}

func NewMetric(host, name string, timestamp uint64, value interface{}, tags map[string]string) *Metric {
	columns := []string{"time", "value", "hostname"}
	points := [][]interface{}{{timestamp, value, host}}
	if tags != nil {
		for k, v := range tags {
			if k == "hostname" {
				points[0][2] = v
			} else {
				columns = append(columns, k)
				points[0] = append(points[0], v)
			}
		}
	}
	return &Metric{name, columns, points}
}

func NewMetricGroup(host, name string, timestamp uint64, values map[string]interface{}, tags map[string]string) *Metric {
	columns := []string{"time", "hostname"}
	points := [][]interface{}{{timestamp, host}}
	for k, v := range values {
		columns = append(columns, k)
		points[0] = append(points[0], v)
	}
	if tags != nil {
		for k, v := range tags {
			if k == "hostname" {
				points[0][1] = v
			} else {
				columns = append(columns, k)
				points[0] = append(points[0], v)
			}
		}
	}
	return &Metric{name, columns, points}
}

func PushMetrics(metrics []*Metric) {
	if len(metrics) == 0 {
		return
	}
	log.Printf("Pushing %d metrics to InfluxDB\n", len(metrics))

	body, err := json.Marshal(metrics)
	if err != nil {
		log.Println(err)
		return
	}

	resp, err := http.Post(dbUrl, "application/json", bytes.NewBuffer(body))
	if err != nil {
		log.Println(err)
	} else if resp.StatusCode != 200 {
		log.Println("Got Response: ", resp.Status)
	}
}

func mapMetrics(data map[string]interface{}) []*Metric {
	metrics := []*Metric{}
	host := data["internalHostname"].(string)
	timestamp := data["collection_timestamp"].(float64)

	delete(data, "internalHostname")
	delete(data, "collection_timestamp")

	values := make(map[string]map[string]interface{})

	log.Printf("Parsing metrics for: %s\n", host)
	for key, value := range data {
		name, ok := rootMetrics[key]
		if ok {
			index := strings.LastIndexAny(name, ".")
			group_name := name[:index]
			group := values[group_name]
			if group == nil {
				group = make(map[string]interface{})
				values[group_name] = group
			}
			group[name[index+1:]] = value
			delete(data, key)
		}
	}

	debug, err := json.MarshalIndent(data, "", "  ")
	if err == nil {
		fmt.Println(string(debug))
	}

	for name, group := range values {
		metrics = append(metrics, NewMetricGroup(host, name, uint64(timestamp*1000), group, nil))
	}
	return metrics
}

func handleIntake(w http.ResponseWriter, req *http.Request) {
	if handleApiKey(w, req) {
		return
	}

	log.Println(req.Method, req.URL.Path)

	body := req.Body
	if req.Header.Get("Content-Encoding") == "deflate" {
		var err error
		body, err = zlib.NewReader(body)
		if err != nil {
			log.Println(err)
			io.WriteString(w, `{"status":"failed"}`)
			return
		}
	}
	buf := bytes.NewBuffer([]byte{})
	_, err := io.Copy(buf, body)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	data := make(map[string]interface{})
	err = json.Unmarshal(buf.Bytes(), &data)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	metrics := mapMetrics(data)
	go PushMetrics(metrics)

	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func handleApi(w http.ResponseWriter, req *http.Request) {
	if handleApiKey(w, req) {
		return
	}

	log.Println(req.Method, req.URL.Path)

	body := req.Body
	if req.Header.Get("Content-Encoding") == "deflate" {
		var err error
		body, err = zlib.NewReader(body)
		if err != nil {
			log.Println(err)
			io.WriteString(w, `{"status":"failed"}`)
			return
		}
	}
	buf := bytes.NewBuffer([]byte{})
	_, err := io.Copy(buf, body)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	} else {
		// log.Println(buf.String())
	}

	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func handleApiKey(w http.ResponseWriter, req *http.Request) bool {
	if req.UserAgent() == "Datadog-Status-Check" {
		io.WriteString(w, "STILL-ALIVE\n")
		return true
	}
	err := req.ParseForm()
	if err != nil {
		log.Println("Error parsing form:", err)
		http.Error(w, err.Error(), 500)
		return true
	}
	values := req.Form
	api_key := values.Get("api_key")
	delete(values, "api_key")
	if len(authApiKey) > 0 && api_key != authApiKey {
		log.Println("Got bad API key:", api_key)
		http.Error(w, "Bad API Key", 403)
		return true
	}
	return false
}

func main() {
	flag.Parse()
	authApiKey = os.Getenv("API_KEY")
	if len(authApiKey) == 0 {
		log.Println("Warning: API_KEY is blank")
	}
	dbUrl = os.Getenv("DB_URL")
	if len(dbUrl) == 0 {
		dbUrl = "http://localhost:8086/db/datadog/series?u=root&p=root"
	}

	log.Println("dd-house listening on", *port)

	http.HandleFunc("/intake", handleIntake)
	http.HandleFunc("/api/v1/series/", handleApi)
	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(*port), nil))
}
<commit_msg>Add some host fields Add disk metrics<commit_after>package main

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
)

var port = flag.Int("port", 8080, "which port to listen on")
var authApiKey string
var dbUrl string

var rootMetrics = map[string]string{
	"agentVersion": "host.agent_version",
	"os":           "host.os",
	"python":       "host.python_version",
	"uuid":         "host.uuid",

	"system.load.1":       "system.load.1",
	"system.load.5":       "system.load.5",
	"system.load.15":      "system.load.15",
	"system.load.norm.1":  "system.load.norm.1",
	"system.load.norm.5":  "system.load.norm.5",
	"system.load.norm.15": "system.load.norm.15",

	"cpuIdle":   "system.cpu.idle",
	"cpuUser":   "system.cpu.user",
	"cpuWait":   "system.cpu.iowait",
	"cpuSystem": "system.cpu.system",
	"cpuStolen": "system.cpu.stolen",

	"memBuffers":       "system.mem.buffered",
	"memPhysPctUsable": "system.mem.pct_usable",
	"memShared":        "system.mem.shared",
	"memPhysTotal":     "system.mem.total",
	"memPhysUsable":    "system.mem.usable",
	"memCached":        "system.mem.cached",
	"memPhysUsed":      "system.mem.used",
	"memPhysFree":      "system.mem.free",

	"memSwapFree":    "system.swap.free",
	"memSwapTotal":   "system.swap.total",
	"memSwapUsed":    "system.swap.used",
	"memSwapPctFree": "system.swap.pct_free",
}

var diskMetrics = []string{
	"device",
	"total",
	"used",
	"free",
	"in_use",
	"mount",
}

type Metric struct {
	Name    string          `json:"name"`
	Columns []string        `json:"columns"`
	Points  [][]interface{} `json:"points"`
}

func NewMetric(host, name string, timestamp uint64, value interface{}, tags map[string]string) *Metric {
	columns := []string{"time", "value", "hostname"}
	points := [][]interface{}{{timestamp, value, host}}
	if tags != nil {
		for k, v := range tags {
			if k == "hostname" {
				points[0][2] = v
			} else {
				columns = append(columns, k)
				points[0] = append(points[0], v)
			}
		}
	}
	return &Metric{name, columns, points}
}

func NewMetricGroup(host, name string, timestamp uint64, values map[string]interface{}, tags map[string]string) *Metric {
	columns := []string{"time", "hostname"}
	points := [][]interface{}{{timestamp, host}}
	for k, v := range values {
		columns = append(columns, k)
		points[0] = append(points[0], v)
	}
	if tags != nil {
		for k, v := range tags {
			if k == "hostname" {
				points[0][1] = v
			} else {
				columns = append(columns, k)
				points[0] = append(points[0], v)
			}
		}
	}
	return &Metric{name, columns, points}
}

func PushMetrics(metrics []*Metric) {
	if len(metrics) == 0 {
		return
	}
	log.Printf("Pushing %d metrics to InfluxDB\n", len(metrics))

	body, err := json.Marshal(metrics)
	if err != nil {
		log.Println(err)
		return
	}

	resp, err := http.Post(dbUrl, "application/json", bytes.NewBuffer(body))
	if err != nil {
		log.Println(err)
	} else if resp.StatusCode != 200 {
		log.Println("Got Response: ", resp.Status)
	}
}

func mapMetrics(data map[string]interface{}) []*Metric {
	metrics := []*Metric{}
	host := data["internalHostname"].(string)
	timestamp := uint64(data["collection_timestamp"].(float64) * 1000)

	delete(data, "apiKey")
	delete(data, "internalHostname")
	delete(data, "collection_timestamp")

	values := make(map[string]map[string]interface{})

	log.Printf("Parsing metrics for: %s\n", host)
	for key, value := range data {
		name, ok := rootMetrics[key]
		if ok {
			index := strings.LastIndexAny(name, ".")
			group_name := name[:index]
			group := values[group_name]
			if group == nil {
				group = make(map[string]interface{})
				values[group_name] = group
			}
			group[name[index+1:]] = value
			delete(data, key)
		}
	}
	for name, group := range values {
		metrics = append(metrics, NewMetricGroup(host, name, timestamp, group, nil))
	}

	metrics = append(metrics, mapDiskMetrics("system.disk", timestamp, host, data["diskUsage"].([]interface{})))
	delete(data, "diskUsage")
	metrics = append(metrics, mapDiskMetrics("system.fs.inodes", timestamp, host, data["inodes"].([]interface{})))
	delete(data, "inodes")

	debug, err := json.MarshalIndent(data, "", "  ")
	if err == nil {
		fmt.Println(string(debug))
	}
	return metrics
}

func mapDiskMetrics(name string, timestamp uint64, host string, data []interface{}) *Metric {
	columns := append([]string{"time", "hostname"}, diskMetrics...)
	points := make([][]interface{}, len(data))
	metric := &Metric{name, columns, points}
	for i, disk := range data {
		fields := disk.([]interface{})
		points[i] = append([]interface{}{timestamp, host}, fields...)
	}
	return metric
}

func handleIntake(w http.ResponseWriter, req *http.Request) {
	if handleApiKey(w, req) {
		return
	}

	log.Println(req.Method, req.URL.Path)

	body := req.Body
	if req.Header.Get("Content-Encoding") == "deflate" {
		var err error
		body, err = zlib.NewReader(body)
		if err != nil {
			log.Println(err)
			io.WriteString(w, `{"status":"failed"}`)
			return
		}
	}
	buf := bytes.NewBuffer([]byte{})
	_, err := io.Copy(buf, body)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	data := make(map[string]interface{})
	err = json.Unmarshal(buf.Bytes(), &data)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	metrics := mapMetrics(data)
	go PushMetrics(metrics)

	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func handleApi(w http.ResponseWriter, req *http.Request) {
	if handleApiKey(w, req) {
		return
	}

	log.Println(req.Method, req.URL.Path)

	body := req.Body
	if req.Header.Get("Content-Encoding") == "deflate" {
		var err error
		body, err = zlib.NewReader(body)
		if err != nil {
			log.Println(err)
			io.WriteString(w, `{"status":"failed"}`)
			return
		}
	}
	buf := bytes.NewBuffer([]byte{})
	_, err := io.Copy(buf, body)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	} else {
		// log.Println(buf.String())
	}

	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func handleApiKey(w http.ResponseWriter, req *http.Request) bool {
	if req.UserAgent() == "Datadog-Status-Check" {
		io.WriteString(w, "STILL-ALIVE\n")
		return true
	}
	err := req.ParseForm()
	if err != nil {
		log.Println("Error parsing form:", err)
		http.Error(w, err.Error(), 500)
		return true
	}
	values := req.Form
	api_key := values.Get("api_key")
	delete(values, "api_key")
	if len(authApiKey) > 0 && api_key != authApiKey {
		log.Println("Got bad API key:", api_key)
		http.Error(w, "Bad API Key", 403)
		return true
	}
	return false
}

func main() {
	flag.Parse()
	authApiKey = os.Getenv("API_KEY")
	if len(authApiKey) == 0 {
		log.Println("Warning: API_KEY is blank")
	}
	dbUrl = os.Getenv("DB_URL")
	if len(dbUrl) == 0 {
		dbUrl = "http://localhost:8086/db/datadog/series?u=root&p=root"
	}

	log.Println("dd-house listening on", *port)

	http.HandleFunc("/intake", handleIntake)
	http.HandleFunc("/api/v1/series/", handleApi)
	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(*port), nil))
}
<|endoftext|>


<commit_before>// Package httpterm provides closable http.Server with extended read timeouts.
package httpterm

import (
	"errors"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// ErrClosing indicating that operation is not allowed as server is closing
var ErrClosing = errors.New("server closing")

// Server traps http.Server, exposes additional fuctionality
type Server struct {
	// HeadReadTimeout defines timeout for reading request headers.
	HeadReadTimeout time.Duration

	// BodyReadTimeout defines timeout for reading request body.
	// This timeout is being applied just before calling the request handler.
	BodyReadTimeout time.Duration

	// IdleTimeout defines for how long connection can be idle between requests.
	IdleTimeout time.Duration

	// NewAsActive controls whether new connection can be idle before issuing
	// a request. By default, new connections will have IdleTimeout allowing for
	// idle period before issuing a request. However, if this flag is set, new
	// connections will be treated as immediately expecting the request, thus
	// HeadReadTimeout will be applied.
	NewAsActive bool

	// CloseOnSignal enables server shutdown on SIGTERM/SIGNINT.
	// Signal handler is registered in Serve() method.
	CloseOnSignal bool

	server   *http.Server // wrapped http server
	listener *rtListener

	lock    sync.Mutex
	closing bool

	// conns is a map of connections which indicates whether connection is active,
	// i.e. there a request being processed (including header handling)
	conns map[net.Conn]bool
}

// NewServer wraps http.Server, which should be already configured.
func NewServer(server *http.Server) *Server {
	return &Server{server: server, conns: make(map[net.Conn]bool)}
}

// Serve behaves as http.Server.Serve on the wrapped server instance
func (s *Server) Serve(l net.Listener) (pending <-chan bool, err error) {
	oldConnState := s.server.ConnState
	newConnState := func(c net.Conn, state http.ConnState) {
		s.updateConnState(c, state)
		// Pass to original handler
		if oldConnState != nil {
			oldConnState(c, state)
		}
	}

	s.server.ConnState = newConnState

	// Wrap with custom listener
	s.listener = &rtListener{
		Listener:    l,
		newAsActive: s.NewAsActive,
		callback:    func(c net.Conn) { newConnState(c, StateData) },
	}

	// Register signal handling for shutdown if requested
	if s.CloseOnSignal {
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt, syscall.SIGTERM)
		go func() {
			<-c
			s.Close()
		}()
	}

	// Serve loop
	err = s.server.Serve(s.listener)

	// Clear error if server closing
	s.lock.Lock()
	if s.closing {
		err = nil
	}
	s.lock.Unlock()

	// Wait for pending requests
	waiter := make(chan bool)
	pending = waiter
	go func() {
		s.listener.wg.Wait()
		close(waiter)
	}()

	return
}

// ListenAndServe listens on the TCP network address from passed http.Server.Addr
// and calls Serve(). Returned channel can be used to wait for pending requests
// and will be closed once everything is handled.
func (s *Server) ListenAndServe() (pending <-chan bool, err error) {
	addr := s.server.Addr
	if addr == "" {
		addr = ":http"
	}
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	return s.Serve(l)
}

// Close server
func (s *Server) Close() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.closing {
		return ErrClosing
	}

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

	s.server.SetKeepAlivesEnabled(false)
	s.closing = true

	// Set a 100ms deadline for all inactive connections (new or idle).
	// If during this period state changes to active, request will be processed
	// with regular request timeout, otherwise connection will be closed.
	deadline := time.Now().Add(100 * time.Millisecond)
	for c, active := range s.conns {
		if !active {
			c.SetReadDeadline(deadline)
		}
	}

	return nil
}

func (s *Server) getTimeout(state http.ConnState) (timeout time.Duration) {
	// Update state for new connection according to policy
	if state == http.StateNew {
		if s.NewAsActive {
			state = StateData
		} else {
			state = http.StateIdle
		}
	}

	switch state {
	case http.StateIdle:
		timeout = s.IdleTimeout

	case StateData:
		timeout = s.HeadReadTimeout

	case http.StateActive:
		timeout = s.BodyReadTimeout
	}

	return
}

func (s *Server) updateConnState(c net.Conn, state http.ConnState) {
	s.lock.Lock()
	defer s.lock.Unlock()

	// Update connection map
	switch state {
	case http.StateNew, http.StateIdle:
		s.conns[c] = false
	case http.StateClosed, http.StateHijacked:
		delete(s.conns, c)
		s.listener.wg.Done()
	case StateData:
		s.conns[c] = true
	}

	if state == http.StateIdle {
		if c, ok := c.(*rtConn); ok {
			c.idle()
		}
	}

	// Update timeout if not closing or new request
	if !s.closing || state == StateData || state == http.StateActive {
		if t := s.getTimeout(state); t != 0 {
			c.SetReadDeadline(time.Now().Add(t))
		}
	}
}

type rtListener struct {
	net.Listener

	newAsActive bool             // set new connections as active
	callback    func(c net.Conn) // data callback

	wg sync.WaitGroup
}

func (l *rtListener) Accept() (c net.Conn, err error) {
	l.wg.Add(1)
	defer func() {
		if c == nil {
			l.wg.Done()
		}
	}()

	c, err = l.Listener.Accept()
	if c != nil {
		c = &rtConn{c, l.newAsActive, l.callback}
	}

	return
}

const (
	// StateData tells when initial request data (header) was received.
	// This varies from http.StateActive, as the latter is issued after headers
	// are parsed.
	StateData http.ConnState = 100 + iota
)

// rtConn is a net.Conn that sets read deadlines for idle and active state.
// It automatically detects requests as first bytes are read after idle state.
type rtConn struct {
	net.Conn

	active   bool             // are we currently processing a request?
	callback func(c net.Conn) // data callback
}

func (c *rtConn) Read(b []byte) (n int, err error) {
	n, err = c.Conn.Read(b)
	if n > 0 && !c.active {
		c.callback(c)
	}
	return
}

func (c *rtConn) idle() {
	c.active = false
}
<commit_msg>Rename StateData to StateHead<commit_after>// Package httpterm provides closable http.Server with extended read timeouts.
package httpterm

import (
	"errors"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

const (
	// StateHead tells when initial request data (header) was received.
	// This varies from http.StateActive, as the latter is issued after headers
	// are parsed.
	StateHead http.ConnState = 100 + iota
)

// ErrClosing indicating that operation is not allowed as server is closing
var ErrClosing = errors.New("server closing")

// Server traps http.Server, exposes additional fuctionality
type Server struct {
	// HeadReadTimeout defines timeout for reading request headers.
	HeadReadTimeout time.Duration

	// BodyReadTimeout defines timeout for reading request body.
	// This timeout is being applied just before calling the request handler.
	BodyReadTimeout time.Duration

	// IdleTimeout defines for how long connection can be idle between requests.
	IdleTimeout time.Duration

	// NewAsActive controls whether new connection can be idle before issuing
	// a request. By default, new connections will have IdleTimeout allowing for
	// idle period before issuing a request. However, if this flag is set, new
	// connections will be treated as immediately expecting the request, thus
	// HeadReadTimeout will be applied.
	NewAsActive bool

	// CloseOnSignal enables server shutdown on SIGTERM/SIGNINT.
	// Signal handler is registered in Serve() method.
	CloseOnSignal bool

	server   *http.Server // wrapped http server
	listener *rtListener

	lock    sync.Mutex
	closing bool

	// conns is a map of connections which indicates whether connection is active,
	// i.e. there a request being processed (including header handling)
	conns map[net.Conn]bool
}

// NewServer wraps http.Server, which should be already configured.
func NewServer(server *http.Server) *Server {
	return &Server{server: server, conns: make(map[net.Conn]bool)}
}

// Serve behaves as http.Server.Serve on the wrapped server instance
func (s *Server) Serve(l net.Listener) (pending <-chan bool, err error) {
	oldConnState := s.server.ConnState
	newConnState := func(c net.Conn, state http.ConnState) {
		s.updateConnState(c, state)
		// Pass to original handler
		if oldConnState != nil {
			oldConnState(c, state)
		}
	}

	s.server.ConnState = newConnState

	// Wrap with custom listener
	s.listener = &rtListener{
		Listener:    l,
		newAsActive: s.NewAsActive,
		callback:    func(c net.Conn) { newConnState(c, StateHead) },
	}

	// Register signal handling for shutdown if requested
	if s.CloseOnSignal {
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt, syscall.SIGTERM)
		go func() {
			<-c
			s.Close()
		}()
	}

	// Serve loop
	err = s.server.Serve(s.listener)

	// Clear error if server closing
	s.lock.Lock()
	if s.closing {
		err = nil
	}
	s.lock.Unlock()

	// Wait for pending requests
	waiter := make(chan bool)
	pending = waiter
	go func() {
		s.listener.wg.Wait()
		close(waiter)
	}()

	return
}

// ListenAndServe listens on the TCP network address from passed http.Server.Addr
// and calls Serve(). Returned channel can be used to wait for pending requests
// and will be closed once everything is handled.
func (s *Server) ListenAndServe() (pending <-chan bool, err error) {
	addr := s.server.Addr
	if addr == "" {
		addr = ":http"
	}
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	return s.Serve(l)
}

// Close server
func (s *Server) Close() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.closing {
		return ErrClosing
	}

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

	s.server.SetKeepAlivesEnabled(false)
	s.closing = true

	// Set a 100ms deadline for all inactive connections (new or idle).
	// If during this period state changes to active, request will be processed
	// with regular request timeout, otherwise connection will be closed.
	deadline := time.Now().Add(100 * time.Millisecond)
	for c, active := range s.conns {
		if !active {
			c.SetReadDeadline(deadline)
		}
	}

	return nil
}

func (s *Server) getTimeout(state http.ConnState) (timeout time.Duration) {
	// Update state for new connection according to policy
	if state == http.StateNew {
		if s.NewAsActive {
			state = StateHead
		} else {
			state = http.StateIdle
		}
	}

	switch state {
	case http.StateIdle:
		timeout = s.IdleTimeout

	case StateHead:
		timeout = s.HeadReadTimeout

	case http.StateActive:
		timeout = s.BodyReadTimeout
	}

	return
}

func (s *Server) updateConnState(c net.Conn, state http.ConnState) {
	s.lock.Lock()
	defer s.lock.Unlock()

	// Update connection map
	switch state {
	case http.StateNew, http.StateIdle:
		s.conns[c] = false
	case http.StateClosed, http.StateHijacked:
		delete(s.conns, c)
		s.listener.wg.Done()
	case StateHead:
		s.conns[c] = true
	}

	if state == http.StateIdle {
		if c, ok := c.(*rtConn); ok {
			c.idle()
		}
	}

	// Update timeout if not closing or new request
	if !s.closing || state == StateHead || state == http.StateActive {
		if t := s.getTimeout(state); t != 0 {
			c.SetReadDeadline(time.Now().Add(t))
		}
	}
}

type rtListener struct {
	net.Listener

	newAsActive bool             // set new connections as active
	callback    func(c net.Conn) // data callback

	wg sync.WaitGroup
}

func (l *rtListener) Accept() (c net.Conn, err error) {
	l.wg.Add(1)
	defer func() {
		if c == nil {
			l.wg.Done()
		}
	}()

	c, err = l.Listener.Accept()
	if c != nil {
		c = &rtConn{c, l.newAsActive, l.callback}
	}

	return
}

// rtConn is a net.Conn that sets read deadlines for idle and active state.
// It automatically detects requests as first bytes are read after idle state.
type rtConn struct {
	net.Conn

	active   bool             // are we currently processing a request?
	callback func(c net.Conn) // data callback
}

func (c *rtConn) Read(b []byte) (n int, err error) {
	n, err = c.Conn.Read(b)
	if n > 0 && !c.active {
		c.callback(c)
	}
	return
}

func (c *rtConn) idle() {
	c.active = false
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, Markover Inc.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.
// Source code and contact info at http://github.com/poptip/ftc

// Package ftc implements an engine.io-compatible server for persistent client-server connections.
package ftc

import (
	"encoding/json"
	"expvar"
	"fmt"
	"net/http"
	"strings"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/golang/glog"
)

// The defaults for options passed to the server.
const (
	DefaultBasePath       = "/engine.io/"
	DefaultCookieName     = "io"
	DefaultPingTimeout    = 60000
	DefaultPingInterval   = 25000
	DefaultUpgradeTimeout = 10000
)

const (
	// Protocol error codes and mappings.
	errorTransportUnknown   = 0
	errorUnknownSID         = 1
	errorBadHandshakeMethod = 2
	errorBadRequest         = 3

	// Query parameters used in client requests.
	paramTransport          = "transport"
	paramSessionID          = "sid"
	paramJSONPResponseIndex = "j"

	// Available transports.
	transportWebSocket = "websocket"
	transportPolling   = "polling"

	// The default time before closed connections are cleaned from
	// the client pool.
	clientReapTimeout = 5 * time.Second
)

var errorMessage = map[int]string{
	errorTransportUnknown:   "Transport unknown",
	errorUnknownSID:         "Session ID unknown",
	errorBadHandshakeMethod: "Bad handshake method",
	errorBadRequest:         "Bad request",
}

var (
	validTransports = map[string]bool{
		transportWebSocket: true,
		transportPolling:   true,
	}
	validUpgrades = map[string]bool{
		transportWebSocket: true,
	}
)

var numClients = expvar.NewInt("num_clients")

// getValidUpgrades returns a slice containing the valid protocols
// that a connection can upgrade to.
func getValidUpgrades() []string {
	upgrades := make([]string, len(validUpgrades))
	i := 0
	for u := range validUpgrades {
		upgrades[i] = u
		i++
	}
	return upgrades
}

// A Handler is called by the server when a connection is
// opened successfully. An example echo handler is shown below.
//   func EchoServer(c *ftc.Conn) {
//   	for {
//   		var msg string
//   		if err := c.Receive(&msg); err != nil {
//   			log.Println("receive: %v", err)
//   			continue
//   		}
//   		if err := c.Send(msg); err != nil {
//   			log.Println("send: %v", err)
//   		}
//   	}
//   }
// It can then be used in combination with NewServer as follows:
//   http.Handle("/engine.io/", ftc.NewServer(nil, ftc.Handler(EchoServer)))
type Handler func(*Conn)

// A server represents a server of an FTC connection.
type server struct {
	// Handler handles an FTC connection.
	Handler

	basePath   string
	cookieName string

	pingTimeout    int
	pingInterval   int
	upgradeTimeout int

	clients  *clientSet        // The set of connections (some may be closed).
	wsServer *websocket.Server // The underlying WebSocket server.
}

// Options are the parameters passed to the server.
type Options struct {
	// BasePath is the base URL path that the server handles requests for.
	BasePath string
	// CookieName is the name of the cookie set upon successful handshake.
	CookieName string
	// DisableCookie is true if no cookies should be set upon handshake.
	DisableCookie bool
	// PingTimeout is how long a ping packet can hang before the Conn is considered closed.
	PingTimeout int
	// PingInterval specifies how often a ping packet should be sent to the server.
	PingInterval int
	// UpgradeTimeout specifies the maximum time an upgrade can take.
	UpgradeTimeout int
}

// NewServer allocates and returns a new server with the given
// options and handler. If nil options are passed, the defaults
// specified in the constants above are used instead.
func NewServer(o *Options, h Handler) *server {
	if o == nil {
		o = &Options{}
	}
	if len(o.BasePath) == 0 {
		o.BasePath = DefaultBasePath
	}
	if len(o.CookieName) == 0 && !o.DisableCookie {
		o.CookieName = DefaultCookieName
	}
	if o.PingInterval == 0 {
		o.PingInterval = DefaultPingInterval
	}
	if o.PingTimeout == 0 {
		o.PingTimeout = DefaultPingTimeout
	}
	if o.UpgradeTimeout == 0 {
		o.UpgradeTimeout = DefaultUpgradeTimeout
	}

	s := &server{
		Handler:        h,
		basePath:       o.BasePath,
		cookieName:     o.CookieName,
		pingTimeout:    o.PingTimeout,
		pingInterval:   o.PingInterval,
		upgradeTimeout: o.UpgradeTimeout,
		clients:        newClientSet(),
	}
	go s.startReaper()
	s.wsServer = &websocket.Server{Handler: s.wsMainHandler}
	return s
}

// startReaper continuously removes closed connections from the
// client set via the reap function.
func (s *server) startReaper() {
	for {
		// TODO(andybons): does this need to be protected by a mutex?
		if s.clients == nil {
			glog.Fatal("server cannot have a nil client set")
		}
		s.clients.reap()
		numClients.Set(int64(s.clients.len()))
		time.Sleep(clientReapTimeout)
	}
}

// receiveWSPacket calls receive on the given WebSocket connection
// and unmarshals it into the given packet.
func receiveWSPacket(ws *websocket.Conn, pkt *packet) error {
	var msg string
	if err := websocket.Message.Receive(ws, &msg); err != nil {
		return err
	}
	glog.Infoln("got websocket message", msg)
	if err := pkt.UnmarshalText([]byte(msg)); err != nil {
		return err
	}
	return nil
}

// sendWSPacket marshals the given packet and sends it over the
// given WebSocket connection.
func sendWSPacket(ws *websocket.Conn, pkt *packet) error {
	if ws == nil || pkt == nil {
		return fmt.Errorf("nil websocket or packet ws: %+v, pkt: %+v", ws, pkt)
	}
	glog.Infof("sending websocket packet: %+v", pkt)
	b, err := pkt.MarshalText()
	if err != nil {
		return err
	}
	if err := websocket.Message.Send(ws, string(b)); err != nil {
		return err
	}
	return nil
}

// handleWSPing send the appropriate response packet (pong) with the
// same data as the given packet over the given WebSocket connection.
func handleWSPing(pkt *packet, ws *websocket.Conn) error {
	resp := newPacket(packetTypePong, pkt.data)
	b, err := resp.MarshalText()
	if err != nil {
		return err
	}
	if err := websocket.Message.Send(ws, string(b)); err != nil {
		return err
	}
	return nil
}

// wsMainHandler continuously receives on the given WebSocket
// connection and delegates the packets received to the appropriate
// handler functions.
func (s *server) wsMainHandler(ws *websocket.Conn) {
	var c *Conn
	var sid string
	for {
		if len(sid) == 0 {
			glog.Infoln("websocket: no session id.")
			sid = ws.Request().FormValue(paramSessionID)
			if len(sid) == 0 {
				sid = s.wsHandshake(ws)
			}
		}
		c = s.clients.get(sid)
		if c == nil {
			glog.Errorln("could not find connection with id", sid)
			continue
		}

		glog.Infoln("waiting for websocket message")
		pkt := &packet{}
		if err := receiveWSPacket(ws, pkt); err != nil {
			glog.Errorln("websocket read message error:", err)
			break
		}
		glog.Infof("got websocket packet %+v from conn %p", pkt, ws)
		if c.upgraded {
			c.onPacket(pkt)
			continue
		}

		// Client has not been upgraded yet.
		if c.ws != nil && c.ws != ws {
			glog.Errorf("websocket already associated with session %s", c.ID)
			break
		}
		c.ws = ws

		switch pkt.typ {
		case packetTypePing:
			if err := handleWSPing(pkt, ws); err != nil {
				glog.Errorf("problem handling websocket ping packet: %+v; error: %v", pkt, err)
				break
			}
			// Force a polling cycle to ensure a fast upgrade.
			if data, _ := pkt.data.(string); data == messageProbe {
				c.sendPacket(newPacket(packetTypeNoop, nil))
			}
		case packetTypeUpgrade:
			c.onUpgrade(pkt)
		}
	}
	glog.Infof("closing websocket connection %p", ws)
	c.close()
}

// pollingHandler handles all XHR polling requests to the server, initiating
// a handshake if the request’s session ID does not already exist within
// the client set.
func (s *server) pollingHandler(w http.ResponseWriter, r *http.Request) {
	sid := r.FormValue(paramSessionID)
	if len(sid) > 0 {
		socket := s.clients.get(sid)
		if socket == nil {
			serverError(w, errorUnknownSID)
			return
		}
		if r.Method == "POST" {
			socket.pollingDataPost(w, r)
			return
		} else if r.Method == "GET" {
			socket.pollingDataGet(w, r)
			return
		}
		http.Error(w, "bad method", http.StatusBadRequest)
		return
	}
	s.pollingHandshake(w, r)
}

// wsHandshake creates a new FTC Conn with the given WebSocket connection
// as the underlying transport and calls the server’s Handler. It returns
// the session ID of the newly-created connection.
func (s *server) wsHandshake(ws *websocket.Conn) string {
	glog.Infof("starting websocket handshake. Handler: %+v", s.Handler)
	c := newConn(s.pingInterval, s.pingTimeout, transportWebSocket)
	s.clients.add(c)
	c.wsOpen(ws)

	go s.Handler(c)
	return c.ID
}

// pollingHandshake creates a new FTC Conn with the given HTTP Request and
// ResponseWriter, setting a persistence cookie if necessary and calling
// the server’s Handler.
func (s *server) pollingHandshake(w http.ResponseWriter, r *http.Request) {
	glog.Infof("starting polling handshake. Handler: %+v", s.Handler)
	c := newConn(s.pingInterval, s.pingTimeout, transportPolling)
	s.clients.add(c)

	if len(s.cookieName) > 0 {
		glog.Infoln("setting cookie:", s.cookieName, c.ID)
		http.SetCookie(w, &http.Cookie{
			Name:  s.cookieName,
			Value: c.ID,
		})
	}
	c.pollingOpen(w, r)
	glog.Infof("polling handleshake complete. calling Handler: %+v", s.Handler)
	go s.Handler(c)
}

// ServeHTTP implements the http.Handler interface for an FTC Server.
func (s *server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	remoteAddr := r.Header.Get("X-Forwarded-For")
	if len(remoteAddr) == 0 {
		remoteAddr = r.RemoteAddr
	}
	glog.Infof("%s (%s) %s %s %s", r.Proto, r.Header.Get("X-Forwarded-Proto"), r.Method, remoteAddr, r.URL)

	transport := r.FormValue(paramTransport)
	if strings.HasPrefix(r.URL.Path, s.basePath) && !validTransports[transport] {
		serverError(w, errorTransportUnknown)
		return
	}

	if transport == transportWebSocket {
		s.wsServer.ServeHTTP(w, r)
		return
	}

	s.pollingHandler(w, r)
}

// serverError sends a JSON-encoded message to the given ResponseWriter
// with the given error code.
func serverError(w http.ResponseWriter, code int) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusBadRequest)
	msg := struct {
		Code    int    `json:"code"`
		Message string `json:"message"`
	}{
		Code:    code,
		Message: errorMessage[code],
	}
	if err := json.NewEncoder(w).Encode(msg); err != nil {
		glog.Errorln("error encoding error msg %+v: %s", msg, err)
		return
	}
}
<commit_msg>printf<commit_after>// Copyright (c) 2014, Markover Inc.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.
// Source code and contact info at http://github.com/poptip/ftc

// Package ftc implements an engine.io-compatible server for persistent client-server connections.
package ftc

import (
	"encoding/json"
	"expvar"
	"fmt"
	"net/http"
	"strings"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/golang/glog"
)

// The defaults for options passed to the server.
const (
	DefaultBasePath       = "/engine.io/"
	DefaultCookieName     = "io"
	DefaultPingTimeout    = 60000
	DefaultPingInterval   = 25000
	DefaultUpgradeTimeout = 10000
)

const (
	// Protocol error codes and mappings.
	errorTransportUnknown   = 0
	errorUnknownSID         = 1
	errorBadHandshakeMethod = 2
	errorBadRequest         = 3

	// Query parameters used in client requests.
	paramTransport          = "transport"
	paramSessionID          = "sid"
	paramJSONPResponseIndex = "j"

	// Available transports.
	transportWebSocket = "websocket"
	transportPolling   = "polling"

	// The default time before closed connections are cleaned from
	// the client pool.
	clientReapTimeout = 5 * time.Second
)

var errorMessage = map[int]string{
	errorTransportUnknown:   "Transport unknown",
	errorUnknownSID:         "Session ID unknown",
	errorBadHandshakeMethod: "Bad handshake method",
	errorBadRequest:         "Bad request",
}

var (
	validTransports = map[string]bool{
		transportWebSocket: true,
		transportPolling:   true,
	}
	validUpgrades = map[string]bool{
		transportWebSocket: true,
	}
)

var numClients = expvar.NewInt("num_clients")

// getValidUpgrades returns a slice containing the valid protocols
// that a connection can upgrade to.
func getValidUpgrades() []string {
	upgrades := make([]string, len(validUpgrades))
	i := 0
	for u := range validUpgrades {
		upgrades[i] = u
		i++
	}
	return upgrades
}

// A Handler is called by the server when a connection is
// opened successfully. An example echo handler is shown below.
//   func EchoServer(c *ftc.Conn) {
//   	for {
//   		var msg string
//   		if err := c.Receive(&msg); err != nil {
//   			log.Printf("receive: %v", err)
//   			continue
//   		}
//   		if err := c.Send(msg); err != nil {
//   			log.Printf("send: %v", err)
//   		}
//   	}
//   }
// It can then be used in combination with NewServer as follows:
//   http.Handle("/engine.io/", ftc.NewServer(nil, ftc.Handler(EchoServer)))
type Handler func(*Conn)

// A server represents a server of an FTC connection.
type server struct {
	// Handler handles an FTC connection.
	Handler

	basePath   string
	cookieName string

	pingTimeout    int
	pingInterval   int
	upgradeTimeout int

	clients  *clientSet        // The set of connections (some may be closed).
	wsServer *websocket.Server // The underlying WebSocket server.
}

// Options are the parameters passed to the server.
type Options struct {
	// BasePath is the base URL path that the server handles requests for.
	BasePath string
	// CookieName is the name of the cookie set upon successful handshake.
	CookieName string
	// DisableCookie is true if no cookies should be set upon handshake.
	DisableCookie bool
	// PingTimeout is how long a ping packet can hang before the Conn is considered closed.
	PingTimeout int
	// PingInterval specifies how often a ping packet should be sent to the server.
	PingInterval int
	// UpgradeTimeout specifies the maximum time an upgrade can take.
	UpgradeTimeout int
}

// NewServer allocates and returns a new server with the given
// options and handler. If nil options are passed, the defaults
// specified in the constants above are used instead.
func NewServer(o *Options, h Handler) *server {
	if o == nil {
		o = &Options{}
	}
	if len(o.BasePath) == 0 {
		o.BasePath = DefaultBasePath
	}
	if len(o.CookieName) == 0 && !o.DisableCookie {
		o.CookieName = DefaultCookieName
	}
	if o.PingInterval == 0 {
		o.PingInterval = DefaultPingInterval
	}
	if o.PingTimeout == 0 {
		o.PingTimeout = DefaultPingTimeout
	}
	if o.UpgradeTimeout == 0 {
		o.UpgradeTimeout = DefaultUpgradeTimeout
	}

	s := &server{
		Handler:        h,
		basePath:       o.BasePath,
		cookieName:     o.CookieName,
		pingTimeout:    o.PingTimeout,
		pingInterval:   o.PingInterval,
		upgradeTimeout: o.UpgradeTimeout,
		clients:        newClientSet(),
	}
	go s.startReaper()
	s.wsServer = &websocket.Server{Handler: s.wsMainHandler}
	return s
}

// startReaper continuously removes closed connections from the
// client set via the reap function.
func (s *server) startReaper() {
	for {
		// TODO(andybons): does this need to be protected by a mutex?
		if s.clients == nil {
			glog.Fatal("server cannot have a nil client set")
		}
		s.clients.reap()
		numClients.Set(int64(s.clients.len()))
		time.Sleep(clientReapTimeout)
	}
}

// receiveWSPacket calls receive on the given WebSocket connection
// and unmarshals it into the given packet.
func receiveWSPacket(ws *websocket.Conn, pkt *packet) error {
	var msg string
	if err := websocket.Message.Receive(ws, &msg); err != nil {
		return err
	}
	glog.Infoln("got websocket message", msg)
	if err := pkt.UnmarshalText([]byte(msg)); err != nil {
		return err
	}
	return nil
}

// sendWSPacket marshals the given packet and sends it over the
// given WebSocket connection.
func sendWSPacket(ws *websocket.Conn, pkt *packet) error {
	if ws == nil || pkt == nil {
		return fmt.Errorf("nil websocket or packet ws: %+v, pkt: %+v", ws, pkt)
	}
	glog.Infof("sending websocket packet: %+v", pkt)
	b, err := pkt.MarshalText()
	if err != nil {
		return err
	}
	if err := websocket.Message.Send(ws, string(b)); err != nil {
		return err
	}
	return nil
}

// handleWSPing send the appropriate response packet (pong) with the
// same data as the given packet over the given WebSocket connection.
func handleWSPing(pkt *packet, ws *websocket.Conn) error {
	resp := newPacket(packetTypePong, pkt.data)
	b, err := resp.MarshalText()
	if err != nil {
		return err
	}
	if err := websocket.Message.Send(ws, string(b)); err != nil {
		return err
	}
	return nil
}

// wsMainHandler continuously receives on the given WebSocket
// connection and delegates the packets received to the appropriate
// handler functions.
func (s *server) wsMainHandler(ws *websocket.Conn) {
	var c *Conn
	var sid string
	for {
		if len(sid) == 0 {
			glog.Infoln("websocket: no session id.")
			sid = ws.Request().FormValue(paramSessionID)
			if len(sid) == 0 {
				sid = s.wsHandshake(ws)
			}
		}
		c = s.clients.get(sid)
		if c == nil {
			glog.Errorln("could not find connection with id", sid)
			continue
		}

		glog.Infoln("waiting for websocket message")
		pkt := &packet{}
		if err := receiveWSPacket(ws, pkt); err != nil {
			glog.Errorln("websocket read message error:", err)
			break
		}
		glog.Infof("got websocket packet %+v from conn %p", pkt, ws)
		if c.upgraded {
			c.onPacket(pkt)
			continue
		}

		// Client has not been upgraded yet.
		if c.ws != nil && c.ws != ws {
			glog.Errorf("websocket already associated with session %s", c.ID)
			break
		}
		c.ws = ws

		switch pkt.typ {
		case packetTypePing:
			if err := handleWSPing(pkt, ws); err != nil {
				glog.Errorf("problem handling websocket ping packet: %+v; error: %v", pkt, err)
				break
			}
			// Force a polling cycle to ensure a fast upgrade.
			if data, _ := pkt.data.(string); data == messageProbe {
				c.sendPacket(newPacket(packetTypeNoop, nil))
			}
		case packetTypeUpgrade:
			c.onUpgrade(pkt)
		}
	}
	glog.Infof("closing websocket connection %p", ws)
	c.close()
}

// pollingHandler handles all XHR polling requests to the server, initiating
// a handshake if the request’s session ID does not already exist within
// the client set.
func (s *server) pollingHandler(w http.ResponseWriter, r *http.Request) {
	sid := r.FormValue(paramSessionID)
	if len(sid) > 0 {
		socket := s.clients.get(sid)
		if socket == nil {
			serverError(w, errorUnknownSID)
			return
		}
		if r.Method == "POST" {
			socket.pollingDataPost(w, r)
			return
		} else if r.Method == "GET" {
			socket.pollingDataGet(w, r)
			return
		}
		http.Error(w, "bad method", http.StatusBadRequest)
		return
	}
	s.pollingHandshake(w, r)
}

// wsHandshake creates a new FTC Conn with the given WebSocket connection
// as the underlying transport and calls the server’s Handler. It returns
// the session ID of the newly-created connection.
func (s *server) wsHandshake(ws *websocket.Conn) string {
	glog.Infof("starting websocket handshake. Handler: %+v", s.Handler)
	c := newConn(s.pingInterval, s.pingTimeout, transportWebSocket)
	s.clients.add(c)
	c.wsOpen(ws)

	go s.Handler(c)
	return c.ID
}

// pollingHandshake creates a new FTC Conn with the given HTTP Request and
// ResponseWriter, setting a persistence cookie if necessary and calling
// the server’s Handler.
func (s *server) pollingHandshake(w http.ResponseWriter, r *http.Request) {
	glog.Infof("starting polling handshake. Handler: %+v", s.Handler)
	c := newConn(s.pingInterval, s.pingTimeout, transportPolling)
	s.clients.add(c)

	if len(s.cookieName) > 0 {
		glog.Infoln("setting cookie:", s.cookieName, c.ID)
		http.SetCookie(w, &http.Cookie{
			Name:  s.cookieName,
			Value: c.ID,
		})
	}
	c.pollingOpen(w, r)
	glog.Infof("polling handleshake complete. calling Handler: %+v", s.Handler)
	go s.Handler(c)
}

// ServeHTTP implements the http.Handler interface for an FTC Server.
func (s *server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	remoteAddr := r.Header.Get("X-Forwarded-For")
	if len(remoteAddr) == 0 {
		remoteAddr = r.RemoteAddr
	}
	glog.Infof("%s (%s) %s %s %s", r.Proto, r.Header.Get("X-Forwarded-Proto"), r.Method, remoteAddr, r.URL)

	transport := r.FormValue(paramTransport)
	if strings.HasPrefix(r.URL.Path, s.basePath) && !validTransports[transport] {
		serverError(w, errorTransportUnknown)
		return
	}

	if transport == transportWebSocket {
		s.wsServer.ServeHTTP(w, r)
		return
	}

	s.pollingHandler(w, r)
}

// serverError sends a JSON-encoded message to the given ResponseWriter
// with the given error code.
func serverError(w http.ResponseWriter, code int) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusBadRequest)
	msg := struct {
		Code    int    `json:"code"`
		Message string `json:"message"`
	}{
		Code:    code,
		Message: errorMessage[code],
	}
	if err := json.NewEncoder(w).Encode(msg); err != nil {
		glog.Errorln("error encoding error msg %+v: %s", msg, err)
		return
	}
}
<|endoftext|>


<commit_before>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"reflect"
	"sync"
)

type Handler func(*Server, *Client, []byte) interface{}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func([]byte) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)
	//Called when no event handler for a specific event exists
	DefaultHandler Handler

	handlers     map[string]Handler
	handlersLock *sync.RWMutex

	newClient chan *Client
	stop      chan struct{}
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string]Handler),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
		stop:      make(chan struct{}),
	}

	go s.listener()
	return s
}

func (s *Server) Close() {
	s.stop <- struct{}{}
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		if r == room {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	s.rooms[r] = append(s.rooms[r], c)
	s.roomsLock.Unlock()

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.RLock()
	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	s.roomsLock.RUnlock()
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		return
	}

	s.roomsLock.Lock()
	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	if len(s.rooms[r]) == 0 {
		delete(s.rooms, r)
	}
	s.roomsLock.Unlock()

	index = -1
	s.joinedRoomsLock.RLock()
	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	s.joinedRoomsLock.RUnlock()
	if index == -1 {
		return
	}

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
	if len(s.joinedRooms[id]) == 0 {
		delete(s.joinedRooms, id)
	}
}

//Send all clients in room room data
func (s *Server) Broadcast(room string, data string) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			c.Emit(data)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		go func(c *Client) {
			c.EmitJSON(v)
		}(client)
	}
	s.roomsLock.RUnlock()
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	rooms := make([]string, len(s.joinedRooms[id]))
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}
func (s *Server) listener() {
	for {
		select {
		case c := <-s.newClient:
			go c.listener(s)
		case <-s.stop:
			return
		}
	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f Handler) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}

//A Receiver interface implements the Name method, which returns a name for the
//event, given a Handler's name
type Receiver interface {
	Name(string) string
}

//Similar to net/rpc's Register, expect that rcvr needs to implement the
//Receiver interface
func (s *Server) Register(rcvr Receiver) {
	rtype := reflect.TypeOf(rcvr)

	for i := 0; i < rtype.NumMethod(); i++ {
		method := rtype.Method(i)
		if method.Name == "Name" {
			continue
		}

		s.On(rcvr.Name(method.Name), func(_ *Server, c *Client, b []byte) interface{} {
			in := []reflect.Value{
				reflect.ValueOf(rcvr),
				reflect.ValueOf(s),
				reflect.ValueOf(c),
				reflect.ValueOf(b)}

			rtrn := method.Func.Call(in)
			return rtrn[0].Interface()
		})

	}
}
<commit_msg>Add Rooms()<commit_after>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"reflect"
	"sync"
)

type Handler func(*Server, *Client, []byte) interface{}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func([]byte) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)
	//Called when no event handler for a specific event exists
	DefaultHandler Handler

	handlers     map[string]Handler
	handlersLock *sync.RWMutex

	newClient chan *Client
	stop      chan struct{}
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string]Handler),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
		stop:      make(chan struct{}),
	}

	go s.listener()
	return s
}

func (s *Server) Close() {
	s.stop <- struct{}{}
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		if r == room {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	s.rooms[r] = append(s.rooms[r], c)
	s.roomsLock.Unlock()

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.RLock()
	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	s.roomsLock.RUnlock()
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		return
	}

	s.roomsLock.Lock()
	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	if len(s.rooms[r]) == 0 {
		delete(s.rooms, r)
	}
	s.roomsLock.Unlock()

	index = -1
	s.joinedRoomsLock.RLock()
	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	s.joinedRoomsLock.RUnlock()
	if index == -1 {
		return
	}

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
	if len(s.joinedRooms[id]) == 0 {
		delete(s.joinedRooms, id)
	}
}

//Send all clients in room room data
func (s *Server) Broadcast(room string, data string) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			c.Emit(data)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		go func(c *Client) {
			c.EmitJSON(v)
		}(client)
	}
	s.roomsLock.RUnlock()
}

//Returns a map of room name -> number of clients
func (s *Server) Rooms() map[string]int {
	rooms := make(map[string]int)

	s.roomsLock.RLock()
	defer s.roomsLock.RUnlock()
	for room, clients := range s.rooms {
		rooms[room] = len(clients)
	}

	return rooms
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	rooms := make([]string, len(s.joinedRooms[id]))
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}
func (s *Server) listener() {
	for {
		select {
		case c := <-s.newClient:
			go c.listener(s)
		case <-s.stop:
			return
		}
	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f Handler) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}

//A Receiver interface implements the Name method, which returns a name for the
//event, given a Handler's name
type Receiver interface {
	Name(string) string
}

//Similar to net/rpc's Register, expect that rcvr needs to implement the
//Receiver interface
func (s *Server) Register(rcvr Receiver) {
	rtype := reflect.TypeOf(rcvr)

	for i := 0; i < rtype.NumMethod(); i++ {
		method := rtype.Method(i)
		if method.Name == "Name" {
			continue
		}

		s.On(rcvr.Name(method.Name), func(_ *Server, c *Client, b []byte) interface{} {
			in := []reflect.Value{
				reflect.ValueOf(rcvr),
				reflect.ValueOf(s),
				reflect.ValueOf(c),
				reflect.ValueOf(b)}

			rtrn := method.Func.Call(in)
			return rtrn[0].Interface()
		})

	}
}
<|endoftext|>


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

package goplay

import (
	"bytes"
	"encoding/json"
	"go/parser"
	"go/printer"
	"go/token"
	"net/http"
)

func init() {
	http.HandleFunc("/fmt", fmtHandler)
}

type fmtResponse struct {
	Body  string
	Error string
}

func fmtHandler(w http.ResponseWriter, r *http.Request) {
	resp := new(fmtResponse)
	body, err := gofmt(r.FormValue("body"))
	if err != nil {
		resp.Error = err.Error()
	} else {
		resp.Body = body
	}
	json.NewEncoder(w).Encode(resp)
}

func gofmt(body string) (string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "prog.go", body, parser.ParseComments)
	if err != nil {
		return "", err
	}
	var buf bytes.Buffer
	err = printer.Fprint(&buf, fset, f)
	if err != nil {
		return "", err
	}
	return buf.String(), nil
}
<commit_msg>goplay: make Format button more consistent with gofmt<commit_after>// Copyright 2012 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package goplay

import (
	"bytes"
	"encoding/json"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"net/http"
)

func init() {
	http.HandleFunc("/fmt", fmtHandler)
}

type fmtResponse struct {
	Body  string
	Error string
}

func fmtHandler(w http.ResponseWriter, r *http.Request) {
	resp := new(fmtResponse)
	body, err := gofmt(r.FormValue("body"))
	if err != nil {
		resp.Error = err.Error()
	} else {
		resp.Body = body
	}
	json.NewEncoder(w).Encode(resp)
}

func gofmt(body string) (string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "prog.go", body, parser.ParseComments)
	if err != nil {
		return "", err
	}
	ast.SortImports(fset, f)
	var buf bytes.Buffer
	err = printer.Fprint(&buf, fset, f)
	if err != nil {
		return "", err
	}
	return buf.String(), nil
}
<|endoftext|>