<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"image"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"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:"image" binding:"required"`
	Timestamp   int64                 `form:"timestamp" binding:"required"`
	Filename    string                `form:"filename"`
	Callback    string                `form:"callback"`
	Signature   string                `form:"signature" binding:"required"`
}

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, secretKey, err := generateKey()
						if err != nil {
							log.Println("Adding a new API key failed, please try again")
							return
						}

						log.Printf("Key added: %s, secret: %s", key, secretKey)
					},
				},
				{
					Name:  "generatesecret",
					Usage: "Generates a new secret key for a given API key (generatesecret [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()
						secret, err := generateSecret(key)
						if err != nil {
							log.Println(err.Error())
							return
						}
						log.Printf("The new secret is: %s. Please save it now.", secret)
					},
				},
				{
					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, make([]*Text, 0)}
	} 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, _ := transformation.createFilePath(baseImagePath)
	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()
}

// UploadResponse is a struct to represent a JSON response for the upload handler
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")
	}

	if uf.PhotoUpload == nil {
		return http.StatusBadRequest, uploadError("missing image field")
	}

	// Check signature only when API key is used
	// Note: when no API key is passed in but required for uploads, the above
	// hasPermission check should fail
	if params["apikey"] != "" {
		uploadTime := time.Unix(uf.Timestamp, 0)
		delta := time.Since(uploadTime).Minutes()
		if delta < 0 || delta > 5 {
			return http.StatusBadRequest, uploadError("invalid timestamp")
		}

		queryParams := make(map[string]string)
		queryParams["timestamp"] = strconv.FormatInt(uf.Timestamp, 10)
		if uf.Callback != "" {
			queryParams["callback"] = uf.Callback
		}
		if uf.Filename != "" {
			queryParams["filename"] = uf.Filename
		}

		secret, err := getSecretForKey(params["apikey"])
		if err != nil {
			return http.StatusBadRequest, uploadError("authorization error")
		}
		if !isValidSignature(uf.Signature, secret, queryParams) {
			return http.StatusBadRequest, uploadError("invalid signature")
		}
	}

	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
	now := time.Now()
	randomInt := rand.Intn(1000)
	baseImagePath := fmt.Sprintf("%d-%d.%s", 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, _ := transformation.createFilePath(baseImagePath)
				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
	}))
}
<commit_msg>Refactored upload code a bit<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"image"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"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:"image" binding:"required"`
	Timestamp   int64                 `form:"timestamp" binding:"required"`
	Signature   string                `form:"signature" binding:"required"`
}

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, secretKey, err := generateKey()
						if err != nil {
							log.Println("Adding a new API key failed, please try again")
							return
						}

						log.Printf("Key added: %s, secret: %s", key, secretKey)
					},
				},
				{
					Name:  "generatesecret",
					Usage: "Generates a new secret key for a given API key (generatesecret [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()
						secret, err := generateSecret(key)
						if err != nil {
							log.Println(err.Error())
							return
						}
						log.Printf("The new secret is: %s. Please save it now.", secret)
					},
				},
				{
					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, make([]*Text, 0)}
	} 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, _ := transformation.createFilePath(baseImagePath)
	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()
}

// UploadResponse is a struct to represent a JSON response for the upload handler
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")
	}

	if uf.PhotoUpload == nil {
		return http.StatusBadRequest, uploadError("missing image field")
	}

	// Check signature only when API key is used
	// Note: when no API key is passed in but required for uploads, the above
	// hasPermission check should fail
	if params["apikey"] != "" {
		uploadTime := time.Unix(uf.Timestamp, 0)
		delta := time.Since(uploadTime).Minutes()
		if delta < 0 || delta > 5 {
			return http.StatusBadRequest, uploadError("invalid timestamp")
		}

		queryParams := make(map[string]string)
		queryParams["timestamp"] = strconv.FormatInt(uf.Timestamp, 10)

		secret, err := getSecretForKey(params["apikey"])
		if err != nil {
			return http.StatusBadRequest, uploadError("authorization error")
		}
		if !isValidSignature(uf.Signature, secret, queryParams) {
			return http.StatusBadRequest, uploadError("invalid signature")
		}
	}

	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
	now := time.Now()
	randomInt := rand.Intn(1000)
	baseImagePath := fmt.Sprintf("%d-%d.%s", 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, _ := transformation.createFilePath(baseImagePath)
				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 (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"log"
	"net"
)

type (
	Header struct {
		ID      uint16
		Flags   uint16
		QDCount uint16
		ANCount uint16
		NSCount uint16
		ARCount uint16
	}
	Query struct {
		QName  string
		QType  QType
		QClass QClass
	}
)

func readHeader(data []byte, pos int) (Header, int) {
	header := Header{}
	header.ID, pos = readUint16(data, pos)
	header.Flags, pos = readUint16(data, pos)
	header.QDCount, pos = readUint16(data, pos)
	header.ANCount, pos = readUint16(data, pos)
	header.NSCount, pos = readUint16(data, pos)
	header.ARCount, pos = readUint16(data, pos)
	return header, pos
}

func readQType(data []byte, pos int) (QType, int) {
	n, pos := readUint16(data, pos)
	return QType(n), pos
}

func readQClass(data []byte, pos int) (QClass, int) {
	n, pos := readUint16(data, pos)
	return QClass(n), pos
}

func readUint16(data []byte, ini int) (uint16, int) {
	fin := ini + 2
	return binary.BigEndian.Uint16(data[ini:fin]), fin
}

func readName(data []byte, pos int) (string, int) {
	var name string
	for data[pos] != 0 {
		labelLength := int(data[pos])
		pos++
		label := data[pos : pos+labelLength]
		pos += labelLength
		name += string(label) + "."
	}
	pos++
	return name, pos
}

func main() {
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4(0, 0, 0, 0), Port: 53})
	if err != nil {
		log.Panicf("server: cannot connect: %s", err)
	}
	defer conn.Close()

	data := make([]byte, 576)
	n, addr, _ := conn.ReadFromUDP(data)
	fmt.Printf("ADDR: %s\n", addr)
	fmt.Printf("RAW DATA: %b\n", data[0:n])

	var pos int
	header, pos := readHeader(data, pos)

	d, _ := json.Marshal(header)
	fmt.Printf("HEADER: %s\n", d)

	q := Query{}
	q.QName, pos = readName(data, pos)
	q.QType, pos = readQType(data, pos)
	q.QClass, pos = readQClass(data, pos)

	d, _ = json.Marshal(q)
	fmt.Printf("MSG: %s\n", d)
}
<commit_msg>(wip) add Msg construct and extract methods<commit_after>package main

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"log"
	"net"
)

type (
	Header struct {
		ID      uint16
		Flags   uint16
		QDCount uint16
		ANCount uint16
		NSCount uint16
		ARCount uint16
	}
	Query struct {
		QName  string
		QType  QType
		QClass QClass
	}
	Msg struct {
		Header  Header
		Queries []Query
	}
)

func parseHeader(data []byte, pos int) (Header, int) {
	header := Header{}
	header.ID, pos = readUint16(data, pos)
	header.Flags, pos = readUint16(data, pos)
	header.QDCount, pos = readUint16(data, pos)
	header.ANCount, pos = readUint16(data, pos)
	header.NSCount, pos = readUint16(data, pos)
	header.ARCount, pos = readUint16(data, pos)

	// TODO(alex): remove.
	d, _ := json.Marshal(header)
	fmt.Printf("HEADER: %s\n", d)

	return header, pos
}

func parseQueries(header Header, data []byte, pos int) ([]Query, int) {
	queries := make([]Query, header.QDCount)
	for i := range queries {
		queries[i].QName, pos = readName(data, pos)
		queries[i].QType, pos = readQType(data, pos)
		queries[i].QClass, pos = readQClass(data, pos)

		// TODO(alex): remove.
		d, _ := json.Marshal(queries[i])
		fmt.Printf("QUERY %d: %s\n", i, d)
	}
	return queries, pos
}

func readQType(data []byte, pos int) (QType, int) {
	n, pos := readUint16(data, pos)
	return QType(n), pos
}

func readQClass(data []byte, pos int) (QClass, int) {
	n, pos := readUint16(data, pos)
	return QClass(n), pos
}

func readUint16(data []byte, ini int) (uint16, int) {
	fin := ini + 2
	return binary.BigEndian.Uint16(data[ini:fin]), fin
}

func readName(data []byte, pos int) (string, int) {
	var name string
	for data[pos] != 0 {
		labelLength := int(data[pos])
		pos++
		label := data[pos : pos+labelLength]
		pos += labelLength
		name += string(label) + "."
	}
	pos++
	return name, pos
}

func main() {
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4(0, 0, 0, 0), Port: 53})
	if err != nil {
		log.Panicf("server: cannot connect: %s", err)
	}
	defer conn.Close()

	data := make([]byte, 576)
	n, addr, _ := conn.ReadFromUDP(data)

	// TODO(alex): remove
	fmt.Printf("ADDR: %s\n", addr)
	fmt.Printf("RAW DATA: %b\n", data[0:n])

	parseMessage(data)
}
func parseMessage(data []byte) Msg {
	var pos int
	msg := Msg{}
	msg.Header, pos = parseHeader(data, pos)
	msg.Queries, pos = parseQueries(msg.Header, data, pos)
	return msg
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/gin-gonic/gin"
	"github.com/rakyll/magicmime"
	"github.com/satori/go.uuid"
	"github.com/spf13/viper"
)

func main() {
	// create gin server
	r := gin.Default()

	// look for config file
	viper.SetConfigName("config")
	viper.AddConfigPath(".")
	err := viper.ReadInConfig()
	// if no config file then read from ENV VARS
	if err != nil {
		viper.AutomaticEnv()
	}
	// setup keys from config file or ENV VARS
	awsID := viper.GetString("AWS_ID")
	awsSecret := viper.GetString("AWS_SECRET")
	awsBucket := viper.GetString("AWS_BUCKET")
	awsRegion := viper.GetString("AWS_REGION")
	awsToken := viper.GetString("AWS_TOKEN")

	// upload route
	r.POST("/upload", func(c *gin.Context) {
		// get file upload
		file, header, err := c.Request.FormFile("file")
		if err != nil {
			log.Println("Failed to upload", err)
			c.JSON(400, gin.H{"error": "you need to provide a file to upload"})
			return
		}
		// get key folder
		folder := c.PostForm("folder")
		// setup s3 uploader
		uploader := s3manager.NewUploader(session.New(&aws.Config{
			Credentials: credentials.NewStaticCredentials(awsID, awsSecret, awsToken),
			Region:      aws.String(awsRegion),
		}))
		// create uuid v4
		u1 := uuid.NewV4()
		// get file extension
		fileExt := filepath.Ext(header.Filename)
		// create unique filename and folder as well
		filename := folder + u1.String() + fileExt
		// upload file to s3
		result, err := uploader.Upload(&s3manager.UploadInput{
			Body:   file,
			Bucket: aws.String(awsBucket),
			Key:    aws.String(filename),
		})
		// if can't download then throw error else return s3 key (url)
		if err != nil {
			log.Println("Failed to upload", err)
			c.JSON(500, gin.H{"error": "there was an error uploading"})
			return
		}
		log.Println("Successfully uploaded to", result.Location)
		c.JSON(200, gin.H{"url": result.Location})
		return
	})

	// download route
	r.POST("/download", func(c *gin.Context) {
		// get AWS key as param
		key := c.PostForm("file")
		folder := c.PostForm("folder")
		if key == "" {
			c.JSON(400, gin.H{"error": "you must provide a file to download"})
			return
		}
		// setup file
		file, err := os.Create(key)
		if err != nil {
			c.JSON(500, gin.H{"error": "there was an error downloading"})
			return
		}
		// close the file and delete after route call is done
		defer file.Close()
		defer os.Remove(key)
		// combine folder and key
		filename := folder + key
		// download file from s3
		downloader := s3manager.NewDownloader(session.New(&aws.Config{
			Credentials: credentials.NewStaticCredentials(awsID, awsSecret, awsToken),
			Region:      aws.String(awsRegion),
		}))
		_, err = downloader.Download(file, &s3.GetObjectInput{
			Bucket: aws.String(awsBucket),
			Key:    aws.String(filename),
		})
		// if can't download from S3
		if err != nil {
			log.Println("Failed to download", err)
			c.JSON(500, gin.H{"error": "there was an error downloading"})
			return
		}
		// init magicmime else throw error
		if err := magicmime.Open(magicmime.MAGIC_MIME_TYPE | magicmime.MAGIC_SYMLINK | magicmime.MAGIC_ERROR); err != nil {
			log.Println("Failed to read mime type", err)
			c.JSON(500, gin.H{"error": "there was an error reading the mime type"})
			return
		}
		// close magicmime after route call is done
		defer magicmime.Close()
		// read file
		bytes, err := ioutil.ReadFile(key)
		// if can't read file then throw error
		if err != nil {
			log.Println("Failed to read file", err)
			c.JSON(500, gin.H{"error": "there was an error opening the file"})
			return
		}
		// read mimetype from file buffer
		mimetype, err := magicmime.TypeByBuffer(bytes)
		// if can't read mimetype then throw error
		if err != nil {
			log.Println("Failed to read mime type", err)
			c.JSON(500, gin.H{"error": "there was an error reading the mime type"})
			return
		}
		// stream data to the requestor
		c.Data(200, mimetype, bytes)
		return
	})

	// run gin server
	r.Run()
}
<commit_msg>change err msg<commit_after>package main

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/gin-gonic/gin"
	"github.com/rakyll/magicmime"
	"github.com/satori/go.uuid"
	"github.com/spf13/viper"
)

func main() {
	// create gin server
	r := gin.Default()

	// look for config file
	viper.SetConfigName("config")
	viper.AddConfigPath(".")
	err := viper.ReadInConfig()
	// if no config file then read from ENV VARS
	if err != nil {
		viper.AutomaticEnv()
	}
	// setup keys from config file or ENV VARS
	awsID := viper.GetString("AWS_ID")
	awsSecret := viper.GetString("AWS_SECRET")
	awsBucket := viper.GetString("AWS_BUCKET")
	awsRegion := viper.GetString("AWS_REGION")
	awsToken := viper.GetString("AWS_TOKEN")

	// upload route
	r.POST("/upload", func(c *gin.Context) {
		// get file upload
		file, header, err := c.Request.FormFile("file")
		if err != nil {
			log.Println("Failed to upload", err)
			c.JSON(400, gin.H{"error": "you need to provide a file to upload"})
			return
		}
		// get key folder
		folder := c.PostForm("folder")
		// setup s3 uploader
		uploader := s3manager.NewUploader(session.New(&aws.Config{
			Credentials: credentials.NewStaticCredentials(awsID, awsSecret, awsToken),
			Region:      aws.String(awsRegion),
		}))
		// create uuid v4
		u1 := uuid.NewV4()
		// get file extension
		fileExt := filepath.Ext(header.Filename)
		// create unique filename and folder as well
		filename := folder + u1.String() + fileExt
		// upload file to s3
		result, err := uploader.Upload(&s3manager.UploadInput{
			Body:   file,
			Bucket: aws.String(awsBucket),
			Key:    aws.String(filename),
		})
		// if can't download then throw error else return s3 key (url)
		if err != nil {
			log.Println("Failed to upload", err)
			c.JSON(500, gin.H{"error": "there was an error uploading"})
			return
		}
		log.Println("Successfully uploaded to", result.Location)
		c.JSON(200, gin.H{"url": result.Location})
		return
	})

	// download route
	r.POST("/download", func(c *gin.Context) {
		// get AWS key as param
		key := c.PostForm("file")
		folder := c.PostForm("folder")
		if key == "" {
			c.JSON(400, gin.H{"error": "you must provide a file to download"})
			return
		}
		// setup file
		file, err := os.Create(key)
		if err != nil {
			c.JSON(500, gin.H{"error": "there was an error creating the tmp file"})
			return
		}
		// close the file and delete after route call is done
		defer file.Close()
		defer os.Remove(key)
		// combine folder and key
		filename := folder + key
		// download file from s3
		downloader := s3manager.NewDownloader(session.New(&aws.Config{
			Credentials: credentials.NewStaticCredentials(awsID, awsSecret, awsToken),
			Region:      aws.String(awsRegion),
		}))
		_, err = downloader.Download(file, &s3.GetObjectInput{
			Bucket: aws.String(awsBucket),
			Key:    aws.String(filename),
		})
		// if can't download from S3
		if err != nil {
			log.Println("Failed to download", err)
			c.JSON(500, gin.H{"error": "there was an error downloading"})
			return
		}
		// init magicmime else throw error
		if err := magicmime.Open(magicmime.MAGIC_MIME_TYPE | magicmime.MAGIC_SYMLINK | magicmime.MAGIC_ERROR); err != nil {
			log.Println("Failed to read mime type", err)
			c.JSON(500, gin.H{"error": "there was an error reading the mime type"})
			return
		}
		// close magicmime after route call is done
		defer magicmime.Close()
		// read file
		bytes, err := ioutil.ReadFile(key)
		// if can't read file then throw error
		if err != nil {
			log.Println("Failed to read file", err)
			c.JSON(500, gin.H{"error": "there was an error opening the file"})
			return
		}
		// read mimetype from file buffer
		mimetype, err := magicmime.TypeByBuffer(bytes)
		// if can't read mimetype then throw error
		if err != nil {
			log.Println("Failed to read mime type", err)
			c.JSON(500, gin.H{"error": "there was an error reading the mime type"})
			return
		}
		// stream data to the requestor
		c.Data(200, mimetype, bytes)
		return
	})

	// run gin server
	r.Run()
}
<|endoftext|>


<commit_before>package veneur

import (
	"bytes"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/Sirupsen/logrus"
	"github.com/getsentry/raven-go"
	"github.com/zenazn/goji/bind"
	"github.com/zenazn/goji/graceful"
)

// must be a var so it can be set at link time
var VERSION = "dirty"

type Server struct {
	Workers     []*Worker
	EventWorker *EventWorker

	statsd *statsd.Client
	logger *logrus.Logger
	sentry *raven.Client

	Hostname string
	Tags     []string

	DDHostname string
	DDAPIKey   string
	HTTPClient *http.Client

	HTTPAddr    string
	ForwardAddr string
	UDPAddr     *net.UDPAddr
	RcvbufBytes int

	HistogramPercentiles []float64
}

func NewFromConfig(conf Config) (ret Server, err error) {
	ret.Hostname = conf.Hostname
	ret.Tags = conf.Tags
	ret.DDHostname = conf.APIHostname
	ret.DDAPIKey = conf.Key
	ret.HistogramPercentiles = conf.Percentiles

	ret.HTTPClient = &http.Client{
		// make sure that POSTs to datadog do not overflow the flush interval
		Timeout: conf.Interval * 9 / 10,
		// we're fine with using the default transport and redirect behavior
	}

	ret.statsd, err = statsd.NewBuffered(conf.StatsAddr, 1024)
	if err != nil {
		return
	}
	ret.statsd.Namespace = "veneur."
	ret.statsd.Tags = append(ret.Tags, "veneurlocalonly")

	// nil is a valid sentry client that noops all methods, if there is no DSN
	// we can just leave it as nil
	if conf.SentryDSN != "" {
		ret.sentry, err = raven.New(conf.SentryDSN)
		if err != nil {
			return
		}
	}

	ret.logger = logrus.New()
	if conf.Debug {
		ret.logger.Level = logrus.DebugLevel
	}
	ret.logger.Hooks.Add(sentryHook{
		c:        ret.sentry,
		hostname: ret.Hostname,
		lv: []logrus.Level{
			logrus.ErrorLevel,
			logrus.FatalLevel,
			logrus.PanicLevel,
		},
	})
	ret.logger.WithField("version", VERSION).Info("Starting server")

	ret.logger.WithField("number", conf.NumWorkers).Info("Starting workers")
	ret.Workers = make([]*Worker, conf.NumWorkers)
	for i := range ret.Workers {
		ret.Workers[i] = NewWorker(i+1, ret.statsd, ret.logger)
		// do not close over loop index
		go func(w *Worker) {
			defer func() {
				ret.ConsumePanic(recover())
			}()
			w.Work()
		}(ret.Workers[i])
	}

	ret.EventWorker = NewEventWorker(ret.statsd)
	go func() {
		defer func() {
			ret.ConsumePanic(recover())
		}()
		ret.EventWorker.Work()
	}()

	ret.UDPAddr, err = net.ResolveUDPAddr("udp", conf.UDPAddr)
	if err != nil {
		return
	}
	ret.RcvbufBytes = conf.ReadBufferSizeBytes
	ret.HTTPAddr = conf.HTTPAddr
	ret.ForwardAddr = conf.ForwardAddr

	conf.Key = "REDACTED"
	conf.SentryDSN = "REDACTED"
	ret.logger.WithField("config", conf).Debug("Initialized server")

	return
}

func (s *Server) HandlePacket(packet []byte) {
	if len(packet) == 0 {
		// a lot of clients send packets that accidentally have a trailing
		// newline, it's easier to just let them be
		return
	}

	if bytes.HasPrefix(packet, []byte{'_', 'e', '{'}) {
		event, err := ParseEvent(packet)
		if err != nil {
			s.logger.WithFields(logrus.Fields{
				logrus.ErrorKey: err,
				"packet":        string(packet),
			}).Error("Could not parse packet")
			s.statsd.Count("packet.error_total", 1, []string{"packet_type:event"}, 1.0)
			return
		}
		s.EventWorker.EventChan <- *event
	} else if bytes.HasPrefix(packet, []byte{'_', 's', 'c'}) {
		svcheck, err := ParseServiceCheck(packet)
		if err != nil {
			s.logger.WithFields(logrus.Fields{
				logrus.ErrorKey: err,
				"packet":        string(packet),
			}).Error("Could not parse packet")
			s.statsd.Count("packet.error_total", 1, []string{"packet_type:service_check"}, 1.0)
			return
		}
		s.EventWorker.ServiceCheckChan <- *svcheck
	} else {
		metric, err := ParseMetric(packet)
		if err != nil {
			s.logger.WithFields(logrus.Fields{
				logrus.ErrorKey: err,
				"packet":        string(packet),
			}).Error("Could not parse packet")
			s.statsd.Count("packet.error_total", 1, []string{"packet_type:metric"}, 1.0)
			return
		}
		s.Workers[metric.Digest%uint32(len(s.Workers))].PacketChan <- *metric
	}
}

func (s *Server) ReadSocket(packetPool *sync.Pool, reuseport bool) {
	// each goroutine gets its own socket
	// if the sockets support SO_REUSEPORT, then this will cause the
	// kernel to distribute datagrams across them, for better read
	// performance
	serverConn, err := NewSocket(s.UDPAddr, s.RcvbufBytes, reuseport)
	if err != nil {
		// if any goroutine fails to create the socket, we can't really
		// recover, so we just blow up
		// this probably indicates a systemic issue, eg lack of
		// SO_REUSEPORT support
		s.logger.WithError(err).Fatal("Error listening for UDP")
	}
	s.logger.WithField("address", s.UDPAddr).Info("UDP server listening")

	for {
		buf := packetPool.Get().([]byte)
		n, _, err := serverConn.ReadFrom(buf)
		if err != nil {
			s.logger.WithError(err).Error("Error reading from UDP")
			continue
		}

		// statsd allows multiple packets to be joined by newlines and sent as
		// one larger packet
		// note that spurious newlines are not allowed in this format, it has
		// to be exactly one newline between each packet, with no leading or
		// trailing newlines
		splitPacket := NewSplitBytes(buf[:n], '\n')
		for splitPacket.Next() {
			s.HandlePacket(splitPacket.Chunk())
		}

		// the Metric struct created by HandlePacket has no byte slices in it,
		// only strings
		// therefore there are no outstanding references to this byte slice, we
		// can return it to the pool
		packetPool.Put(buf)
	}
}

func (s *Server) HTTPServe() {
	httpSocket := bind.Socket(s.HTTPAddr)
	graceful.Timeout(10 * time.Second)
	graceful.PreHook(func() {
		s.logger.Info("Terminating HTTP listener")
	})
	graceful.HandleSignals()
	s.logger.WithField("address", s.HTTPAddr).Info("HTTP server listening")
	bind.Ready()

	if err := graceful.Serve(httpSocket, s.Handler()); err != nil {
		s.logger.WithError(err).Error("HTTP server shut down due to error")
	}
	graceful.Wait()
}

// SplitBytes iterates over a byte buffer, returning chunks split by a given
// delimiter byte. It does not perform any allocations, and does not modify the
// buffer it is given. It is not safe for use by concurrent goroutines.
//
//     sb := NewSplitBytes(buf, '\n')
//     for sb.Next() {
//         fmt.Printf("%q\n", sb.Chunk())
//     }
//
// The sequence of chunks returned by SplitBytes is equivalent to calling
// bytes.Split, except without allocating an intermediate slice.
type SplitBytes struct {
	buf          []byte
	delim        byte
	currentChunk []byte
	lastChunk    bool
}

func NewSplitBytes(buf []byte, delim byte) *SplitBytes {
	return &SplitBytes{
		buf:   buf,
		delim: delim,
	}
}

// Next advances SplitBytes to the next chunk, returning true if a new chunk
// actually exists and false otherwise.
func (sb *SplitBytes) Next() bool {
	if sb.lastChunk {
		// we do not check the length here, this ensures that we return the
		// last chunk in the sequence (even if it's empty)
		return false
	}

	next := bytes.IndexByte(sb.buf, sb.delim)
	if next == -1 {
		// no newline, consume the entire buffer
		sb.currentChunk = sb.buf
		sb.buf = nil
		sb.lastChunk = true
	} else {
		sb.currentChunk = sb.buf[:next]
		sb.buf = sb.buf[next+1:]
	}
	return true
}

// Chunk returns the current chunk.
func (sb *SplitBytes) Chunk() []byte {
	return sb.currentChunk
}
<commit_msg>Add documentation to server.go<commit_after>package veneur

import (
	"bytes"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/Sirupsen/logrus"
	"github.com/getsentry/raven-go"
	"github.com/zenazn/goji/bind"
	"github.com/zenazn/goji/graceful"
)

// VERSION stores the current veneur version.
// It must be a var so it can be set at link time.
var VERSION = "dirty"

// A Server is the actual veneur instance that will be run.
type Server struct {
	Workers     []*Worker
	EventWorker *EventWorker

	statsd *statsd.Client
	logger *logrus.Logger
	sentry *raven.Client

	Hostname string
	Tags     []string

	DDHostname string
	DDAPIKey   string
	HTTPClient *http.Client

	HTTPAddr    string
	ForwardAddr string
	UDPAddr     *net.UDPAddr
	RcvbufBytes int

	HistogramPercentiles []float64
}

// NewFromConfig creates a new veneur server from a configuration specification.
func NewFromConfig(conf Config) (ret Server, err error) {
	ret.Hostname = conf.Hostname
	ret.Tags = conf.Tags
	ret.DDHostname = conf.APIHostname
	ret.DDAPIKey = conf.Key
	ret.HistogramPercentiles = conf.Percentiles

	ret.HTTPClient = &http.Client{
		// make sure that POSTs to datadog do not overflow the flush interval
		Timeout: conf.Interval * 9 / 10,
		// we're fine with using the default transport and redirect behavior
	}

	ret.statsd, err = statsd.NewBuffered(conf.StatsAddr, 1024)
	if err != nil {
		return
	}
	ret.statsd.Namespace = "veneur."
	ret.statsd.Tags = append(ret.Tags, "veneurlocalonly")

	// nil is a valid sentry client that noops all methods, if there is no DSN
	// we can just leave it as nil
	if conf.SentryDSN != "" {
		ret.sentry, err = raven.New(conf.SentryDSN)
		if err != nil {
			return
		}
	}

	ret.logger = logrus.New()
	if conf.Debug {
		ret.logger.Level = logrus.DebugLevel
	}
	ret.logger.Hooks.Add(sentryHook{
		c:        ret.sentry,
		hostname: ret.Hostname,
		lv: []logrus.Level{
			logrus.ErrorLevel,
			logrus.FatalLevel,
			logrus.PanicLevel,
		},
	})
	ret.logger.WithField("version", VERSION).Info("Starting server")

	ret.logger.WithField("number", conf.NumWorkers).Info("Starting workers")
	ret.Workers = make([]*Worker, conf.NumWorkers)
	for i := range ret.Workers {
		ret.Workers[i] = NewWorker(i+1, ret.statsd, ret.logger)
		// do not close over loop index
		go func(w *Worker) {
			defer func() {
				ret.ConsumePanic(recover())
			}()
			w.Work()
		}(ret.Workers[i])
	}

	ret.EventWorker = NewEventWorker(ret.statsd)
	go func() {
		defer func() {
			ret.ConsumePanic(recover())
		}()
		ret.EventWorker.Work()
	}()

	ret.UDPAddr, err = net.ResolveUDPAddr("udp", conf.UDPAddr)
	if err != nil {
		return
	}
	ret.RcvbufBytes = conf.ReadBufferSizeBytes
	ret.HTTPAddr = conf.HTTPAddr
	ret.ForwardAddr = conf.ForwardAddr

	conf.Key = "REDACTED"
	conf.SentryDSN = "REDACTED"
	ret.logger.WithField("config", conf).Debug("Initialized server")

	return
}

// HandlePacket processes each packet that is sent to the server, and sends to an
// appropriate worker (EventWorker or Worker).
func (s *Server) HandlePacket(packet []byte) {
	if len(packet) == 0 {
		// a lot of clients send packets that accidentally have a trailing
		// newline, it's easier to just let them be
		return
	}

	if bytes.HasPrefix(packet, []byte{'_', 'e', '{'}) {
		event, err := ParseEvent(packet)
		if err != nil {
			s.logger.WithFields(logrus.Fields{
				logrus.ErrorKey: err,
				"packet":        string(packet),
			}).Error("Could not parse packet")
			s.statsd.Count("packet.error_total", 1, []string{"packet_type:event"}, 1.0)
			return
		}
		s.EventWorker.EventChan <- *event
	} else if bytes.HasPrefix(packet, []byte{'_', 's', 'c'}) {
		svcheck, err := ParseServiceCheck(packet)
		if err != nil {
			s.logger.WithFields(logrus.Fields{
				logrus.ErrorKey: err,
				"packet":        string(packet),
			}).Error("Could not parse packet")
			s.statsd.Count("packet.error_total", 1, []string{"packet_type:service_check"}, 1.0)
			return
		}
		s.EventWorker.ServiceCheckChan <- *svcheck
	} else {
		metric, err := ParseMetric(packet)
		if err != nil {
			s.logger.WithFields(logrus.Fields{
				logrus.ErrorKey: err,
				"packet":        string(packet),
			}).Error("Could not parse packet")
			s.statsd.Count("packet.error_total", 1, []string{"packet_type:metric"}, 1.0)
			return
		}
		s.Workers[metric.Digest%uint32(len(s.Workers))].PacketChan <- *metric
	}
}

// ReadSocket listens for available packets to handle.
func (s *Server) ReadSocket(packetPool *sync.Pool, reuseport bool) {
	// each goroutine gets its own socket
	// if the sockets support SO_REUSEPORT, then this will cause the
	// kernel to distribute datagrams across them, for better read
	// performance
	serverConn, err := NewSocket(s.UDPAddr, s.RcvbufBytes, reuseport)
	if err != nil {
		// if any goroutine fails to create the socket, we can't really
		// recover, so we just blow up
		// this probably indicates a systemic issue, eg lack of
		// SO_REUSEPORT support
		s.logger.WithError(err).Fatal("Error listening for UDP")
	}
	s.logger.WithField("address", s.UDPAddr).Info("UDP server listening")

	for {
		buf := packetPool.Get().([]byte)
		n, _, err := serverConn.ReadFrom(buf)
		if err != nil {
			s.logger.WithError(err).Error("Error reading from UDP")
			continue
		}

		// statsd allows multiple packets to be joined by newlines and sent as
		// one larger packet
		// note that spurious newlines are not allowed in this format, it has
		// to be exactly one newline between each packet, with no leading or
		// trailing newlines
		splitPacket := NewSplitBytes(buf[:n], '\n')
		for splitPacket.Next() {
			s.HandlePacket(splitPacket.Chunk())
		}

		// the Metric struct created by HandlePacket has no byte slices in it,
		// only strings
		// therefore there are no outstanding references to this byte slice, we
		// can return it to the pool
		packetPool.Put(buf)
	}
}

// HTTPServe starts the HTTP server and listens perpetually until it encounters an unrecoverable error.
func (s *Server) HTTPServe() {
	httpSocket := bind.Socket(s.HTTPAddr)
	graceful.Timeout(10 * time.Second)
	graceful.PreHook(func() {
		s.logger.Info("Terminating HTTP listener")
	})
	graceful.HandleSignals()
	s.logger.WithField("address", s.HTTPAddr).Info("HTTP server listening")
	bind.Ready()

	if err := graceful.Serve(httpSocket, s.Handler()); err != nil {
		s.logger.WithError(err).Error("HTTP server shut down due to error")
	}
	graceful.Wait()
}

// SplitBytes iterates over a byte buffer, returning chunks split by a given
// delimiter byte. It does not perform any allocations, and does not modify the
// buffer it is given. It is not safe for use by concurrent goroutines.
//
//     sb := NewSplitBytes(buf, '\n')
//     for sb.Next() {
//         fmt.Printf("%q\n", sb.Chunk())
//     }
//
// The sequence of chunks returned by SplitBytes is equivalent to calling
// bytes.Split, except without allocating an intermediate slice.
type SplitBytes struct {
	buf          []byte
	delim        byte
	currentChunk []byte
	lastChunk    bool
}

// NewSplitBytes initializes a SplitBytes struct with the provided buffer and delimiter.
func NewSplitBytes(buf []byte, delim byte) *SplitBytes {
	return &SplitBytes{
		buf:   buf,
		delim: delim,
	}
}

// Next advances SplitBytes to the next chunk, returning true if a new chunk
// actually exists and false otherwise.
func (sb *SplitBytes) Next() bool {
	if sb.lastChunk {
		// we do not check the length here, this ensures that we return the
		// last chunk in the sequence (even if it's empty)
		return false
	}

	next := bytes.IndexByte(sb.buf, sb.delim)
	if next == -1 {
		// no newline, consume the entire buffer
		sb.currentChunk = sb.buf
		sb.buf = nil
		sb.lastChunk = true
	} else {
		sb.currentChunk = sb.buf[:next]
		sb.buf = sb.buf[next+1:]
	}
	return true
}

// Chunk returns the current chunk.
func (sb *SplitBytes) Chunk() []byte {
	return sb.currentChunk
}
<|endoftext|>


<commit_before>// Package mannersagain combines manners and goagain to provide graceful hot
// restarting of net/http servers.
package mannersagain

import (
	"errors"
	"log"
	"net"
	"net/http"
	"time"

	"github.com/braintree/manners"
	"github.com/rcrowley/goagain"
)

func newListener(l net.Listener) net.Listener {
	return listener{Listener: l, closed: make(chan struct{})}
}

type listener struct {
	net.Listener
	closed chan struct{}
}

var ErrClosed = errors.New("mannersagain: listener has been gracefully closed")

func (l listener) Accept() (net.Conn, error) {
	for {
		select {
		case <-l.closed:
			return nil, ErrClosed
		default:
		}

		// Set a deadline so Accept doesn't block forever, which gives
		// us an opportunity to stop gracefully.
		l.Listener.(*net.TCPListener).SetDeadline(time.Now().Add(100e6))
		c, err := l.Listener.Accept()
		if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
			continue
		}
		return c, err
	}
}

func (l listener) Close() error {
	close(l.closed)
	return nil
}

func ListenAndServe(addr string, handler http.Handler) error {
	var gl *manners.GracefulListener
	srv := manners.NewServer()

	done := make(chan struct{})
	serve := func(l net.Listener) {
		srv.Serve(l, handler)
		close(done)
	}

	// Attempt to inherit a listener from our parent
	l, err := goagain.Listener()
	if err != nil {
		// We don't have an inherited listener, create a new one
		l, err = net.Listen("tcp", addr)
		if err != nil {
			return err
		}
		log.Println("Listening on", l.Addr())
		gl = manners.NewListener(newListener(l), srv)
		go serve(gl)
	} else {
		log.Println("Resuming listening on", l.Addr())
		gl = manners.NewListener(newListener(l), srv)
		go serve(gl)

		// If this is the child, send the parent SIGUSR2. If this is the
		// parent, send the child SIGQUIT.
		if err := goagain.Kill(); nil != err {
			return err
		}
	}

	// Block the main goroutine awaiting signals.
	sig, err := goagain.Wait(l)
	if err != nil {
		return err
	}

	// Stop accepting new connections
	gl.Close()
	// Wait for all existing connections to complete
	<-done

	if goagain.Strategy == goagain.Double {
		// If we received SIGUSR2, re-exec the parent process.
		if sig == goagain.SIGUSR2 {
			if err := goagain.Exec(l); err != nil {
				return err
			}
		}
	}

	return nil
}
<commit_msg>Use goagain fork for zombie fix<commit_after>// Package mannersagain combines manners and goagain to provide graceful hot
// restarting of net/http servers.
package mannersagain

import (
	"errors"
	"log"
	"net"
	"net/http"
	"time"

	"github.com/braintree/manners"
	"github.com/titanous/goagain"
)

func newListener(l net.Listener) net.Listener {
	return listener{Listener: l, closed: make(chan struct{})}
}

type listener struct {
	net.Listener
	closed chan struct{}
}

var ErrClosed = errors.New("mannersagain: listener has been gracefully closed")

func (l listener) Accept() (net.Conn, error) {
	for {
		select {
		case <-l.closed:
			return nil, ErrClosed
		default:
		}

		// Set a deadline so Accept doesn't block forever, which gives
		// us an opportunity to stop gracefully.
		l.Listener.(*net.TCPListener).SetDeadline(time.Now().Add(100e6))
		c, err := l.Listener.Accept()
		if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
			continue
		}
		return c, err
	}
}

func (l listener) Close() error {
	close(l.closed)
	return nil
}

func ListenAndServe(addr string, handler http.Handler) error {
	var gl *manners.GracefulListener
	srv := manners.NewServer()

	done := make(chan struct{})
	serve := func(l net.Listener) {
		srv.Serve(l, handler)
		close(done)
	}

	// Attempt to inherit a listener from our parent
	l, err := goagain.Listener()
	if err != nil {
		// We don't have an inherited listener, create a new one
		l, err = net.Listen("tcp", addr)
		if err != nil {
			return err
		}
		log.Println("Listening on", l.Addr())
		gl = manners.NewListener(newListener(l), srv)
		go serve(gl)
	} else {
		log.Println("Resuming listening on", l.Addr())
		gl = manners.NewListener(newListener(l), srv)
		go serve(gl)

		if err := goagain.Kill(); nil != err {
			return err
		}
	}

	// Block the main goroutine awaiting signals.
	sig, err := goagain.Wait(l)
	if err != nil {
		return err
	}

	// Stop accepting new connections
	gl.Close()
	// Wait for all existing connections to complete
	<-done

	if goagain.Strategy == goagain.Double {
		// If we received SIGUSR2, re-exec the parent process.
		if sig == goagain.SIGUSR2 {
			if err := goagain.Exec(l); err != nil {
				return err
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Jacob Taylor jacob@ablox.io
// License: Apache2 - http://www.apache.org/licenses/LICENSE-2.0
package main

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"sort"
	"time"
)

// Event stores the relevant information on events or updates to the storage layer.
type Event struct {
	Source        string
	Name          string
	Message       string
	Time          time.Time
	IsDirectory   bool
	NetworkSource string
}

var events = make([]Event, 0, 100)

// FileEvent stores the nodeid, string, and time of file related events
type FileEvent struct {
	NodeID int32
	Name   string
	Time   time.Time
}

// SendEvent gets events that have happened off to the peer servers so they can replicate the same change
func SendEvent(event Event, fullPath string) {
	// sendEvent to manager
	sendEvent(&event, fullPath, globalSettings.ManagerAddress, globalSettings.ManagerCredentials)

	// SendEvent to all peers
	for k, v := range serverMap {
		fmt.Printf("Considering sending to: %s\n", k)
		if k != globalSettings.Name {
			fmt.Printf("sending to peer %s at %s\n", k, v.Address)
			sendEvent(&event, fullPath, v.Address, globalSettings.ManagerCredentials)
		}
	}
}

func sendEvent(event *Event, fullPath string, address string, credentials string) {
	url := "http://" + address + "/event/"
	fmt.Printf("target url: %s\n", url)

	// Set the event source  (server name)
	event.Source = globalSettings.Name

	jsonStr, _ := json.Marshal(event)
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonStr))
	req.Header.Set("Content-Type", "application/json")

	data := []byte(credentials)
	authHash := base64.StdEncoding.EncodeToString(data)
	req.Header.Add("Authorization", "Basic "+authHash)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		log.Println(err)
		return
	}

	defer resp.Body.Close()

	fmt.Println("response Status:", resp.Status)
	fmt.Println("response Headers:", resp.Header)
	body, _ := ioutil.ReadAll(resp.Body)
	fmt.Println("response Body:", string(body))

	if event.Name == "notify.Write" {
		fmt.Printf("source: notify.Write => %s\n", fullPath)
		url := "http://" + address + "/upload/"
		postFile(event.Message, fullPath, url, credentials)
	}
}

func eventHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		json.NewEncoder(w).Encode(events)
	case "POST":
		log.Println("Got POST: ", r.Body)
		decoder := json.NewDecoder(r.Body)
		var event Event
		err := decoder.Decode(&event)
		if err != nil {
			panic("bad json body")
		}

		log.Println(event.Name + ", path: " + event.Message)
		log.Printf("Event info: %v\n", event)

		pathName := globalSettings.Directory + "/" + event.Message
		relativePath := event.Message

		//todo remember these events and skip the logging on the other side. Possibly use nodeID?

		switch event.Name {
		case "notify.Create":
			fmt.Printf("notify.Create: %s\n", pathName)

			server, exists := serverMap[globalSettings.Name]
			if !exists {
				panic("Unable to find server definition")
			}

			fmt.Println("eventHandler->CreatePath")
			server.storage.CreatePath(relativePath, event.IsDirectory)
			fmt.Println("eventHandler->/CreatePath")
		case "notify.Remove":
			err = os.Remove(pathName)
			if err != nil && !os.IsNotExist(err) {
				panic(fmt.Sprintf("Error deleting folder %s: %v\n", pathName, err))
			}
			fmt.Printf("notify.Remove: %s\n", pathName)
		// todo fix this to handle the two rename events to be one event
		case "notify.Rename":
			fmt.Printf("notify.Rename: %s\n", pathName)

			server, exists := serverMap[globalSettings.Name]
			if !exists {
				panic("Unable to find server definition")
			}

			fmt.Println("eventHandler->CreatePath")
			server.storage.CreatePath(relativePath, event.IsDirectory)
			fmt.Println("eventHandler->/CreatePath")

		default:
			fmt.Printf("Unknown event found, doing nothing. Event: %v\n", event)
		}

		events = append([]Event{event}, events...)
	}
}

/*
Send the folder tree from this node to another node for comparison
*/
func sendFolderTree(initialTree DirTreeMap) {
	var tree = make(DirTreeMap)
	prefixLength := 1 // the slash prefix
	for key, value := range initialTree {
		if key == "" {
			continue
		}
		fmt.Printf("key is: '%s', length is: %d\n", key, prefixLength)
		key = key[prefixLength:]
		tree[key] = value
	}

	// package up the tree
	jsonStr, err := json.Marshal(tree)
	if err != nil {
		panic(err)
	}
	buffer := bytes.NewBuffer(jsonStr)
	client := &http.Client{}

	data := []byte(globalSettings.ManagerCredentials)
	authHash := base64.StdEncoding.EncodeToString(data)

	for k, v := range serverMap {
		// don't send an update to ourselves
		if k == globalSettings.Name {
			continue
		}

		url := "http://" + v.Address + "/tree/"
		fmt.Printf("Posting folder tree to node: %s at URL: %s\n", k, url)

		req, err := http.NewRequest("POST", url, buffer)
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Basic "+authHash)

		resp, err := client.Do(req)
		if err != nil {
			fmt.Printf("we encountered an error!\n%s\n", err)
			//panic(err)
			continue
		}

		//fmt.Println("response Status:", resp.Status)
		//fmt.Println("response Headers:", resp.Header)
		_, _ = ioutil.ReadAll(resp.Body)
		//body, _ := ioutil.ReadAll(resp.Body)
		//fmt.Println("response Body:", string(body))
		resp.Body.Close()
	}
}

func getListOfFolders(w http.ResponseWriter) {
	listOfFileInfo, err := scanDirectoryContents()
	if err != nil {
		log.Fatal(err)
	}
	json.NewEncoder(w).Encode(listOfFileInfo)
	fmt.Printf("Sending tree of size %d back to client\n", len(listOfFileInfo))
}

func deletePaths(deletedPaths []string) {
	if globalSettings.Directory == "" {
		panic("globalSettings.Directory is not configured correctly. Aborting")
	}

	fmt.Println("Paths were deleted on the other side. Delete them here")

	fmt.Printf("Paths to delete\nBefore sort:\n%v\n", deletedPaths)
	// Reverse sort the paths so the most specific is first. This allows us to get away without a recursive delete
	sort.Sort(sort.Reverse(sort.StringSlice(deletedPaths)))
	fmt.Printf("Paths to delete after sort\nAfter sort:\n%v\n", deletedPaths)
	for _, relativePath := range deletedPaths {
		if relativePath == "" || relativePath == "/" {
			fmt.Printf("We had a request to delete the base path. Skipping: %s\n", relativePath)
			continue
		}
		fullPath := globalSettings.Directory + relativePath
		fmt.Printf("Full path is: %s\n", fullPath)

		fmt.Printf("%s: about to remove\n", fullPath)

		// stop on any error except for not exist. We are trying to delete it anyway (or rather, it should have been deleted already)
		err := os.Remove(fullPath)
		if err != nil && !os.IsNotExist(err) {
			panic(err)
		}
		fmt.Printf("%s: done removing (err = %v)\n", fullPath, err)
	}
}

func addPaths(newPaths []string) {
	fmt.Println("Paths were added on the other side. Create them here")
	for _, newPathName := range newPaths {
		fmt.Printf("pathname is: %s\n", newPathName)
		err := os.Mkdir(newPathName, os.ModeDir+os.ModePerm)
		if err != nil && !os.IsExist(err) {
			panic(err)
		}
	}
}

func folderTreeHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		getListOfFolders(w)
	case "POST":
		decoder := json.NewDecoder(r.Body)

		var remoteTreePreTranslate DirTreeMap
		err := decoder.Decode(&remoteTreePreTranslate)
		if err != nil {
			panic("bad json body")
		}

		var remoteTree = make(DirTreeMap)
		for key, value := range remoteTreePreTranslate {
			key = fmt.Sprintf("%s/%s", globalSettings.Directory, key)
			remoteTree[key] = value
		}

		// todo get the node name in here
		fmt.Printf("Received a tree map from: %s\n%s\n", r.RemoteAddr, remoteTree)

		// let's compare to the current one we have.
		_, _, newPaths, deletedPaths, _ := checkForChanges(remoteTree, nil)

		// update the paths to be consistent. Remember, everything is backwards since this is comparing the other side to us instead of two successive states
		// deletion will show up as a new path.

		fmt.Printf("about to check for deletions: %s\n", newPaths)
		// delete folders that were deleted. We delete first, then add to make sure that the old ones will not be in the way
		if len(newPaths) > 0 {
			deletePaths(newPaths)
		}

		// add folders that were created
		fmt.Printf("about to check for new folders: %s\n", deletedPaths)
		// add folders that were added.
		if len(deletedPaths) > 0 {
			addPaths(deletedPaths)
		}
	}
}

func postFile(filename string, fullPath string, targetUrl string, credentials string) error {
	body := &bytes.Buffer{}
	bodyWriter := multipart.NewWriter(body)
	fileWriter, err := bodyWriter.CreateFormFile("uploadfile", filename)
	if err != nil {
		fmt.Println("error writing to buffer")
		return err
	}

	file, err := os.Open(fullPath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(fileWriter, file)
	if err != nil {
		fmt.Println("error copying file")
		return err
	}

	myHash, err := fileMd5Hash(fullPath)
	if err != nil {
		fmt.Printf("failed to calculate MD5 Hash for %s\n", fullPath)
		return err
	}

	bodyWriter.WriteField("HASH", myHash)
	contentType := bodyWriter.FormDataContentType()
	err = bodyWriter.Close()
	if err != nil {
		return err
	}

	req, err := http.NewRequest("POST", targetUrl, body)
	req.Header.Set("Content-Type", contentType)

	data := []byte(credentials)
	authHash := base64.StdEncoding.EncodeToString(data)
	req.Header.Add("Authorization", "Basic "+authHash)

	client := &http.Client{}
	resp, err := client.Do(req)
	defer resp.Body.Close()

	if err != nil {
		return err
	}

	return nil
}

func fileMd5Hash(filePath string) (string, error) {
	var returnMD5String string
	file, err := os.Open(filePath)
	if err != nil {
		return returnMD5String, err
	}
	defer file.Close()

	hash := md5.New()
	_, err = io.Copy(hash, file)
	if err != nil {
		return returnMD5String, err
	}

	hashInBytes := hash.Sum(nil)[:16]

	returnMD5String = hex.EncodeToString(hashInBytes)

	return returnMD5String, nil

}
<commit_msg>removing unused data structure<commit_after>// Copyright 2016 Jacob Taylor jacob@ablox.io
// License: Apache2 - http://www.apache.org/licenses/LICENSE-2.0
package main

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"sort"
	"time"
)

// Event stores the relevant information on events or updates to the storage layer.
type Event struct {
	Source        string
	Name          string
	Message       string
	Time          time.Time
	IsDirectory   bool
	NetworkSource string
}

var events = make([]Event, 0, 100)

// SendEvent gets events that have happened off to the peer servers so they can replicate the same change
func SendEvent(event Event, fullPath string) {

	// look back through the events for a similar event in the recent path.

	// sendEvent to manager
	sendEvent(&event, fullPath, globalSettings.ManagerAddress, globalSettings.ManagerCredentials)

	// SendEvent to all peers
	for k, v := range serverMap {
		fmt.Printf("Considering sending to: %s\n", k)
		if k != globalSettings.Name {
			fmt.Printf("sending to peer %s at %s\n", k, v.Address)
			sendEvent(&event, fullPath, v.Address, globalSettings.ManagerCredentials)
		}
	}
}

func sendEvent(event *Event, fullPath string, address string, credentials string) {
	url := "http://" + address + "/event/"
	fmt.Printf("target url: %s\n", url)

	// Set the event source  (server name)
	event.Source = globalSettings.Name

	jsonStr, _ := json.Marshal(event)
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonStr))
	req.Header.Set("Content-Type", "application/json")

	data := []byte(credentials)
	authHash := base64.StdEncoding.EncodeToString(data)
	req.Header.Add("Authorization", "Basic "+authHash)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		log.Println(err)
		return
	}

	defer resp.Body.Close()

	fmt.Println("response Status:", resp.Status)
	fmt.Println("response Headers:", resp.Header)
	body, _ := ioutil.ReadAll(resp.Body)
	fmt.Println("response Body:", string(body))

	if event.Name == "notify.Write" {
		fmt.Printf("source: notify.Write => %s\n", fullPath)
		url := "http://" + address + "/upload/"
		postFile(event.Message, fullPath, url, credentials)
	}
}

func eventHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		json.NewEncoder(w).Encode(events)
	case "POST":
		log.Println("Got POST: ", r.Body)
		decoder := json.NewDecoder(r.Body)
		var event Event
		err := decoder.Decode(&event)
		if err != nil {
			panic("bad json body")
		}

		log.Println(event.Name + ", path: " + event.Message)
		log.Printf("Event info: %v\n", event)

		pathName := globalSettings.Directory + "/" + event.Message
		relativePath := event.Message

		//todo remember these events and skip the logging on the other side. Possibly use nodeID?

		switch event.Name {
		case "notify.Create":
			fmt.Printf("notify.Create: %s\n", pathName)

			server, exists := serverMap[globalSettings.Name]
			if !exists {
				panic("Unable to find server definition")
			}

			fmt.Println("eventHandler->CreatePath")
			server.storage.CreatePath(relativePath, event.IsDirectory)
			fmt.Println("eventHandler->/CreatePath")
		case "notify.Remove":
			err = os.Remove(pathName)
			if err != nil && !os.IsNotExist(err) {
				panic(fmt.Sprintf("Error deleting folder %s: %v\n", pathName, err))
			}
			fmt.Printf("notify.Remove: %s\n", pathName)
		// todo fix this to handle the two rename events to be one event
		case "notify.Rename":
			fmt.Printf("notify.Rename: %s\n", pathName)

			server, exists := serverMap[globalSettings.Name]
			if !exists {
				panic("Unable to find server definition")
			}

			fmt.Println("eventHandler->CreatePath")
			server.storage.CreatePath(relativePath, event.IsDirectory)
			fmt.Println("eventHandler->/CreatePath")

		default:
			fmt.Printf("Unknown event found, doing nothing. Event: %v\n", event)
		}

		events = append([]Event{event}, events...)
	}
}

/*
Send the folder tree from this node to another node for comparison
*/
func sendFolderTree(initialTree DirTreeMap) {
	var tree = make(DirTreeMap)
	prefixLength := 1 // the slash prefix
	for key, value := range initialTree {
		if key == "" {
			continue
		}
		fmt.Printf("key is: '%s', length is: %d\n", key, prefixLength)
		key = key[prefixLength:]
		tree[key] = value
	}

	// package up the tree
	jsonStr, err := json.Marshal(tree)
	if err != nil {
		panic(err)
	}
	buffer := bytes.NewBuffer(jsonStr)
	client := &http.Client{}

	data := []byte(globalSettings.ManagerCredentials)
	authHash := base64.StdEncoding.EncodeToString(data)

	for k, v := range serverMap {
		// don't send an update to ourselves
		if k == globalSettings.Name {
			continue
		}

		url := "http://" + v.Address + "/tree/"
		fmt.Printf("Posting folder tree to node: %s at URL: %s\n", k, url)

		req, err := http.NewRequest("POST", url, buffer)
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Basic "+authHash)

		resp, err := client.Do(req)
		if err != nil {
			fmt.Printf("we encountered an error!\n%s\n", err)
			//panic(err)
			continue
		}

		//fmt.Println("response Status:", resp.Status)
		//fmt.Println("response Headers:", resp.Header)
		_, _ = ioutil.ReadAll(resp.Body)
		//body, _ := ioutil.ReadAll(resp.Body)
		//fmt.Println("response Body:", string(body))
		resp.Body.Close()
	}
}

func getListOfFolders(w http.ResponseWriter) {
	listOfFileInfo, err := scanDirectoryContents()
	if err != nil {
		log.Fatal(err)
	}
	json.NewEncoder(w).Encode(listOfFileInfo)
	fmt.Printf("Sending tree of size %d back to client\n", len(listOfFileInfo))
}

func deletePaths(deletedPaths []string) {
	if globalSettings.Directory == "" {
		panic("globalSettings.Directory is not configured correctly. Aborting")
	}

	fmt.Println("Paths were deleted on the other side. Delete them here")

	fmt.Printf("Paths to delete\nBefore sort:\n%v\n", deletedPaths)
	// Reverse sort the paths so the most specific is first. This allows us to get away without a recursive delete
	sort.Sort(sort.Reverse(sort.StringSlice(deletedPaths)))
	fmt.Printf("Paths to delete after sort\nAfter sort:\n%v\n", deletedPaths)
	for _, relativePath := range deletedPaths {
		if relativePath == "" || relativePath == "/" {
			fmt.Printf("We had a request to delete the base path. Skipping: %s\n", relativePath)
			continue
		}
		fullPath := globalSettings.Directory + relativePath
		fmt.Printf("Full path is: %s\n", fullPath)

		fmt.Printf("%s: about to remove\n", fullPath)

		// stop on any error except for not exist. We are trying to delete it anyway (or rather, it should have been deleted already)
		err := os.Remove(fullPath)
		if err != nil && !os.IsNotExist(err) {
			panic(err)
		}
		fmt.Printf("%s: done removing (err = %v)\n", fullPath, err)
	}
}

func addPaths(newPaths []string) {
	fmt.Println("Paths were added on the other side. Create them here")
	for _, newPathName := range newPaths {
		fmt.Printf("pathname is: %s\n", newPathName)
		err := os.Mkdir(newPathName, os.ModeDir+os.ModePerm)
		if err != nil && !os.IsExist(err) {
			panic(err)
		}
	}
}

func folderTreeHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		getListOfFolders(w)
	case "POST":
		decoder := json.NewDecoder(r.Body)

		var remoteTreePreTranslate DirTreeMap
		err := decoder.Decode(&remoteTreePreTranslate)
		if err != nil {
			panic("bad json body")
		}

		var remoteTree = make(DirTreeMap)
		for key, value := range remoteTreePreTranslate {
			key = fmt.Sprintf("%s/%s", globalSettings.Directory, key)
			remoteTree[key] = value
		}

		// todo get the node name in here
		fmt.Printf("Received a tree map from: %s\n%s\n", r.RemoteAddr, remoteTree)

		// let's compare to the current one we have.
		_, _, newPaths, deletedPaths, _ := checkForChanges(remoteTree, nil)

		// update the paths to be consistent. Remember, everything is backwards since this is comparing the other side to us instead of two successive states
		// deletion will show up as a new path.

		fmt.Printf("about to check for deletions: %s\n", newPaths)
		// delete folders that were deleted. We delete first, then add to make sure that the old ones will not be in the way
		if len(newPaths) > 0 {
			deletePaths(newPaths)
		}

		// add folders that were created
		fmt.Printf("about to check for new folders: %s\n", deletedPaths)
		// add folders that were added.
		if len(deletedPaths) > 0 {
			addPaths(deletedPaths)
		}
	}
}

func postFile(filename string, fullPath string, targetUrl string, credentials string) error {
	body := &bytes.Buffer{}
	bodyWriter := multipart.NewWriter(body)
	fileWriter, err := bodyWriter.CreateFormFile("uploadfile", filename)
	if err != nil {
		fmt.Println("error writing to buffer")
		return err
	}

	file, err := os.Open(fullPath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(fileWriter, file)
	if err != nil {
		fmt.Println("error copying file")
		return err
	}

	myHash, err := fileMd5Hash(fullPath)
	if err != nil {
		fmt.Printf("failed to calculate MD5 Hash for %s\n", fullPath)
		return err
	}

	bodyWriter.WriteField("HASH", myHash)
	contentType := bodyWriter.FormDataContentType()
	err = bodyWriter.Close()
	if err != nil {
		return err
	}

	req, err := http.NewRequest("POST", targetUrl, body)
	req.Header.Set("Content-Type", contentType)

	data := []byte(credentials)
	authHash := base64.StdEncoding.EncodeToString(data)
	req.Header.Add("Authorization", "Basic "+authHash)

	client := &http.Client{}
	resp, err := client.Do(req)
	defer resp.Body.Close()

	if err != nil {
		return err
	}

	return nil
}

func fileMd5Hash(filePath string) (string, error) {
	var returnMD5String string
	file, err := os.Open(filePath)
	if err != nil {
		return returnMD5String, err
	}
	defer file.Close()

	hash := md5.New()
	_, err = io.Copy(hash, file)
	if err != nil {
		return returnMD5String, err
	}

	hashInBytes := hash.Sum(nil)[:16]

	returnMD5String = hex.EncodeToString(hashInBytes)

	return returnMD5String, nil

}
<|endoftext|>


<commit_before>package milter

import (
	"net"
)

/* RunServer provides a convinient way to start milter server */
func RunServer(server net.Listener, MilterInit func() (Milter, uint32, uint32)) error {
	for {
		// accept connection from client
		client, err := server.Accept()
		if err != nil {
			return err
		}
		// create milter object
		milter, actions, protocol := MilterInit()
		session := MilterSession{
			Actions:  actions,
			Protocol: protocol,
			Sock:     client,
			Milter:   milter,
		}
		// handle connection commands
		go session.HandleMilterCommands()
	}
}
<commit_msg>define MilterInit function type for reference<commit_after>package milter

import (
	"net"
)

/* MilterInit is a function that initializes milter object options */
type MilterInit func() (Milter, uint32, uint32)

/* RunServer provides a convinient way to start milter server */
func RunServer(server net.Listener, init MilterInit) error {
	for {
		// accept connection from client
		client, err := server.Accept()
		if err != nil {
			return err
		}
		// create milter object
		milter, actions, protocol := init()
		session := MilterSession{
			Actions:  actions,
			Protocol: protocol,
			Sock:     client,
			Milter:   milter,
		}
		// handle connection commands
		go session.HandleMilterCommands()
	}
}
<|endoftext|>


<commit_before>// +build linux

package fs2

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/opencontainers/runc/libcontainer/cgroups"
	"github.com/opencontainers/runc/libcontainer/configs"
	"github.com/pkg/errors"
)

type manager struct {
	config *configs.Cgroup
	// dirPath is like "/sys/fs/cgroup/user.slice/user-1001.slice/session-1.scope"
	dirPath string
	// controllers is content of "cgroup.controllers" file.
	// excludes pseudo-controllers ("devices" and "freezer").
	controllers map[string]struct{}
	rootless    bool
}

// NewManager creates a manager for cgroup v2 unified hierarchy.
// dirPath is like "/sys/fs/cgroup/user.slice/user-1001.slice/session-1.scope".
// If dirPath is empty, it is automatically set using config.
func NewManager(config *configs.Cgroup, dirPath string, rootless bool) (cgroups.Manager, error) {
	if config == nil {
		config = &configs.Cgroup{}
	}
	if dirPath == "" {
		var err error
		dirPath, err = defaultDirPath(config)
		if err != nil {
			return nil, err
		}
	}

	m := &manager{
		config:   config,
		dirPath:  dirPath,
		rootless: rootless,
	}
	return m, nil
}

func (m *manager) getControllers() error {
	if m.controllers != nil {
		return nil
	}

	file := filepath.Join(m.dirPath, "cgroup.controllers")
	data, err := ioutil.ReadFile(file)
	if err != nil && !m.rootless {
		return err
	}
	fields := strings.Fields(string(data))
	m.controllers = make(map[string]struct{}, len(fields))
	for _, c := range fields {
		m.controllers[c] = struct{}{}
	}

	return nil
}

func (m *manager) Apply(pid int) error {
	if err := CreateCgroupPath(m.dirPath, m.config); err != nil {
		return err
	}
	if err := cgroups.WriteCgroupProc(m.dirPath, pid); err != nil && !m.rootless {
		return err
	}
	return nil
}

func (m *manager) GetPids() ([]int, error) {
	return cgroups.GetPids(m.dirPath)
}

func (m *manager) GetAllPids() ([]int, error) {
	return cgroups.GetAllPids(m.dirPath)
}

func (m *manager) GetStats() (*cgroups.Stats, error) {
	var (
		errs []error
	)

	st := cgroups.NewStats()
	if err := m.getControllers(); err != nil {
		return st, err
	}

	// pids (since kernel 4.5)
	if _, ok := m.controllers["pids"]; ok {
		if err := statPids(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	} else {
		if err := statPidsWithoutController(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	// memory (since kernel 4.5)
	if _, ok := m.controllers["memory"]; ok {
		if err := statMemory(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	// io (since kernel 4.5)
	if _, ok := m.controllers["io"]; ok {
		if err := statIo(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	// cpu (since kernel 4.15)
	if _, ok := m.controllers["cpu"]; ok {
		if err := statCpu(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	// hugetlb (since kernel 5.6)
	if _, ok := m.controllers["hugetlb"]; ok {
		if err := statHugeTlb(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	if len(errs) > 0 && !m.rootless {
		return st, errors.Errorf("error while statting cgroup v2: %+v", errs)
	}
	return st, nil
}

func (m *manager) Freeze(state configs.FreezerState) error {
	if err := setFreezer(m.dirPath, state); err != nil {
		return err
	}
	m.config.Resources.Freezer = state
	return nil
}

func (m *manager) Destroy() error {
	return os.RemoveAll(m.dirPath)
}

// GetPaths is for compatibility purpose and should be removed in future
func (m *manager) GetPaths() map[string]string {
	_ = m.getControllers()
	paths := map[string]string{
		// pseudo-controller for compatibility
		"devices": m.dirPath,
		"freezer": m.dirPath,
	}
	for c := range m.controllers {
		paths[c] = m.dirPath
	}
	return paths
}

func (m *manager) GetUnifiedPath() (string, error) {
	return m.dirPath, nil
}

func (m *manager) Set(container *configs.Config) error {
	if container == nil || container.Cgroups == nil {
		return nil
	}
	if err := m.getControllers(); err != nil {
		return err
	}
	var errs []error
	// pids (since kernel 4.5)
	if _, ok := m.controllers["pids"]; ok {
		if err := setPids(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// memory (since kernel 4.5)
	if _, ok := m.controllers["memory"]; ok {
		if err := setMemory(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// io (since kernel 4.5)
	if _, ok := m.controllers["io"]; ok {
		if err := setIo(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// cpu (since kernel 4.15)
	if _, ok := m.controllers["cpu"]; ok {
		if err := setCpu(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// devices (since kernel 4.15, pseudo-controller)
	if err := setDevices(m.dirPath, container.Cgroups); err != nil {
		errs = append(errs, err)
	}
	// cpuset (since kernel 5.0)
	if _, ok := m.controllers["cpuset"]; ok {
		if err := setCpuset(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// hugetlb (since kernel 5.6)
	if _, ok := m.controllers["hugetlb"]; ok {
		if err := setHugeTlb(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// freezer (since kernel 5.2, pseudo-controller)
	if err := setFreezer(m.dirPath, container.Cgroups.Freezer); err != nil {
		errs = append(errs, err)
	}
	if len(errs) > 0 && !m.rootless {
		return errors.Errorf("error while setting cgroup v2: %+v", errs)
	}
	m.config = container.Cgroups
	return nil
}

func (m *manager) GetCgroups() (*configs.Cgroup, error) {
	return m.config, nil
}
<commit_msg>cgroups/fs2/Destroy: use Remove, ignore ENOENT<commit_after>// +build linux

package fs2

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/opencontainers/runc/libcontainer/cgroups"
	"github.com/opencontainers/runc/libcontainer/configs"
	"github.com/pkg/errors"
)

type manager struct {
	config *configs.Cgroup
	// dirPath is like "/sys/fs/cgroup/user.slice/user-1001.slice/session-1.scope"
	dirPath string
	// controllers is content of "cgroup.controllers" file.
	// excludes pseudo-controllers ("devices" and "freezer").
	controllers map[string]struct{}
	rootless    bool
}

// NewManager creates a manager for cgroup v2 unified hierarchy.
// dirPath is like "/sys/fs/cgroup/user.slice/user-1001.slice/session-1.scope".
// If dirPath is empty, it is automatically set using config.
func NewManager(config *configs.Cgroup, dirPath string, rootless bool) (cgroups.Manager, error) {
	if config == nil {
		config = &configs.Cgroup{}
	}
	if dirPath == "" {
		var err error
		dirPath, err = defaultDirPath(config)
		if err != nil {
			return nil, err
		}
	}

	m := &manager{
		config:   config,
		dirPath:  dirPath,
		rootless: rootless,
	}
	return m, nil
}

func (m *manager) getControllers() error {
	if m.controllers != nil {
		return nil
	}

	file := filepath.Join(m.dirPath, "cgroup.controllers")
	data, err := ioutil.ReadFile(file)
	if err != nil && !m.rootless {
		return err
	}
	fields := strings.Fields(string(data))
	m.controllers = make(map[string]struct{}, len(fields))
	for _, c := range fields {
		m.controllers[c] = struct{}{}
	}

	return nil
}

func (m *manager) Apply(pid int) error {
	if err := CreateCgroupPath(m.dirPath, m.config); err != nil {
		return err
	}
	if err := cgroups.WriteCgroupProc(m.dirPath, pid); err != nil && !m.rootless {
		return err
	}
	return nil
}

func (m *manager) GetPids() ([]int, error) {
	return cgroups.GetPids(m.dirPath)
}

func (m *manager) GetAllPids() ([]int, error) {
	return cgroups.GetAllPids(m.dirPath)
}

func (m *manager) GetStats() (*cgroups.Stats, error) {
	var (
		errs []error
	)

	st := cgroups.NewStats()
	if err := m.getControllers(); err != nil {
		return st, err
	}

	// pids (since kernel 4.5)
	if _, ok := m.controllers["pids"]; ok {
		if err := statPids(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	} else {
		if err := statPidsWithoutController(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	// memory (since kernel 4.5)
	if _, ok := m.controllers["memory"]; ok {
		if err := statMemory(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	// io (since kernel 4.5)
	if _, ok := m.controllers["io"]; ok {
		if err := statIo(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	// cpu (since kernel 4.15)
	if _, ok := m.controllers["cpu"]; ok {
		if err := statCpu(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	// hugetlb (since kernel 5.6)
	if _, ok := m.controllers["hugetlb"]; ok {
		if err := statHugeTlb(m.dirPath, st); err != nil {
			errs = append(errs, err)
		}
	}
	if len(errs) > 0 && !m.rootless {
		return st, errors.Errorf("error while statting cgroup v2: %+v", errs)
	}
	return st, nil
}

func (m *manager) Freeze(state configs.FreezerState) error {
	if err := setFreezer(m.dirPath, state); err != nil {
		return err
	}
	m.config.Resources.Freezer = state
	return nil
}

func (m *manager) Destroy() error {
	if err := os.Remove(m.dirPath); err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}

// GetPaths is for compatibility purpose and should be removed in future
func (m *manager) GetPaths() map[string]string {
	_ = m.getControllers()
	paths := map[string]string{
		// pseudo-controller for compatibility
		"devices": m.dirPath,
		"freezer": m.dirPath,
	}
	for c := range m.controllers {
		paths[c] = m.dirPath
	}
	return paths
}

func (m *manager) GetUnifiedPath() (string, error) {
	return m.dirPath, nil
}

func (m *manager) Set(container *configs.Config) error {
	if container == nil || container.Cgroups == nil {
		return nil
	}
	if err := m.getControllers(); err != nil {
		return err
	}
	var errs []error
	// pids (since kernel 4.5)
	if _, ok := m.controllers["pids"]; ok {
		if err := setPids(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// memory (since kernel 4.5)
	if _, ok := m.controllers["memory"]; ok {
		if err := setMemory(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// io (since kernel 4.5)
	if _, ok := m.controllers["io"]; ok {
		if err := setIo(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// cpu (since kernel 4.15)
	if _, ok := m.controllers["cpu"]; ok {
		if err := setCpu(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// devices (since kernel 4.15, pseudo-controller)
	if err := setDevices(m.dirPath, container.Cgroups); err != nil {
		errs = append(errs, err)
	}
	// cpuset (since kernel 5.0)
	if _, ok := m.controllers["cpuset"]; ok {
		if err := setCpuset(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// hugetlb (since kernel 5.6)
	if _, ok := m.controllers["hugetlb"]; ok {
		if err := setHugeTlb(m.dirPath, container.Cgroups); err != nil {
			errs = append(errs, err)
		}
	}
	// freezer (since kernel 5.2, pseudo-controller)
	if err := setFreezer(m.dirPath, container.Cgroups.Freezer); err != nil {
		errs = append(errs, err)
	}
	if len(errs) > 0 && !m.rootless {
		return errors.Errorf("error while setting cgroup v2: %+v", errs)
	}
	m.config = container.Cgroups
	return nil
}

func (m *manager) GetCgroups() (*configs.Cgroup, error) {
	return m.config, nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charmstore

import (
	"fmt"
	"net/http"
	"sort"

	"labix.org/v2/mgo"

	"github.com/juju/charmstore/internal/charmstore"
	"github.com/juju/charmstore/internal/v4"
)

// Versions of the API that can be served.
const (
	V4 = "v4"
)

var versions = map[string]func(*charmstore.Store) http.Handler{
	V4: v4.New,
}

// Versions returns all known API versions.
func Versions() []string {
	vs := make([]string, 0, len(versions))
	for v := range versions {
		vs = append(vs, v)
	}
	sort.Strings(vs)
	return vs
}

// NewServer returns a new handler that handles
// charm store requests and stores its data in the given database.
func NewServer(db *mgo.Database, serveVersions ...string) (http.Handler, error) {
	newAPIs := make(map[string]charmstore.NewAPIHandler)
	for _, vers := range serveVersions {
		newAPI := versions[vers]
		if newAPI == nil {
			return nil, fmt.Errorf("unknown version %q", vers)
		}
		newAPIs[vers] = newAPI
	}

	return charmstore.NewServer(db, newAPIs)
}
<commit_msg>charmstore: fix doc comment for review<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charmstore

import (
	"fmt"
	"net/http"
	"sort"

	"labix.org/v2/mgo"

	"github.com/juju/charmstore/internal/charmstore"
	"github.com/juju/charmstore/internal/v4"
)

// Versions of the API that can be served.
const (
	V4 = "v4"
)

var versions = map[string]func(*charmstore.Store) http.Handler{
	V4: v4.New,
}

// Versions returns all known API version strings in alphabetical order.
func Versions() []string {
	vs := make([]string, 0, len(versions))
	for v := range versions {
		vs = append(vs, v)
	}
	sort.Strings(vs)
	return vs
}

// NewServer returns a new handler that handles
// charm store requests and stores its data in the given database.
func NewServer(db *mgo.Database, serveVersions ...string) (http.Handler, error) {
	newAPIs := make(map[string]charmstore.NewAPIHandler)
	for _, vers := range serveVersions {
		newAPI := versions[vers]
		if newAPI == nil {
			return nil, fmt.Errorf("unknown version %q", vers)
		}
		newAPIs[vers] = newAPI
	}

	return charmstore.NewServer(db, newAPIs)
}
<|endoftext|>


<commit_before>package main

import (
  "os"
  "io/ioutil"
  "github.com/go-martini/martini"
  "github.com/martini-contrib/render"
  "code.google.com/p/go.crypto/ssh"
  "time"
  "runtime"
  "encoding/json"
  "log"
)

type EndpointConfiguration struct {
  User string
  Host string
  Command string
}

type AppConfiguration struct {
  PrivateKeyFile string
  AuthMethods []ssh.AuthMethod
  Source EndpointConfiguration
  Destination EndpointConfiguration
}

type AppState struct {
  Working bool
  StartTime *time.Time
  Configuration AppConfiguration
  Logger *log.Logger
}

func NewAppState() *AppState {
  file, err := os.Open("config.json")
  if err != nil {
    panic("Could not open config.json: " + err.Error())
  }

  decoder := json.NewDecoder(file)

  config := AppConfiguration{}
  if err := decoder.Decode(&config); err != nil {
    panic("Could not decode config.json: " + err.Error())
  }

  privKeyText, err := ioutil.ReadFile(config.PrivateKeyFile)
  if err != nil {
    panic("Failed to read private key: " + err.Error())
  }

  privKey, err := ssh.ParseRawPrivateKey(privKeyText)
  if err != nil {
    panic("Failed to parse private key: " + err.Error())
  }

  signer, err := ssh.NewSignerFromKey(privKey)
  if err != nil {
    panic("Failed to create signer: " + err.Error())
  }

  config.AuthMethods = []ssh.AuthMethod{
    ssh.PublicKeys(signer),
  }

  logger := log.New(os.Stdout, "[big-red] ", 0)

  return &AppState{false, nil, config, logger}
}

func main() {
  state := NewAppState()

  m := martini.Classic()
  m.Use(render.Renderer())

  m.Get("/", func(r render.Render) {
    r.HTML(200, "index", state)
  })

  m.Get("/press", func(r render.Render) {
    if !state.Working {
      go state.PerformDump()
    }

    r.Redirect("/")
  })

  m.Run()
}

func (state *AppState) PerformDump() {
  defer func() {
    if r := recover(); r != nil {
      buf := make([]byte, 1<<16)
      runtime.Stack(buf, false)
      state.Logger.Println(r)
      state.Logger.Println(buf)
    }

    state.Logger.Println("Done. Took " + state.Elapsed().String())

    state.Working = false
    state.StartTime = nil
  }()

  start := time.Now()
  state.Working = true
  state.StartTime = &start

  state.Logger.Println("Started performing work")

  sourceSession := state.newSourceSession()
  defer sourceSession.Close()

  destSession := state.newDestinationSession()
  defer destSession.Close()

  stdoutPipe, err := sourceSession.StdoutPipe()
  if err != nil {
    panic("Could create pipe: " + err.Error())
  }

  destSession.Stdin  = stdoutPipe
  destSession.Stdout = os.Stdout
  destSession.Stderr = os.Stderr

  sourceSession.Stderr = os.Stderr

  if err := destSession.Start(state.Configuration.Destination.Command); err != nil {
      panic("Failed to run destination command: " + err.Error())
  }

  if err := sourceSession.Start(state.Configuration.Source.Command); err != nil {
      panic("Failed to run source command: " + err.Error())
  }

  sourceSession.Wait()
  destSession.Wait()
}

func (state *AppState) Elapsed() time.Duration {
  return time.Since(*state.StartTime)
}

func (state *AppState) newSourceSession() *ssh.Session {
  return state.newSession(state.Configuration.Source.User, state.Configuration.Source.Host)
}

func (state *AppState) newDestinationSession() *ssh.Session {
  return state.newSession(state.Configuration.Destination.User, state.Configuration.Destination.Host)
}

func (state *AppState) newSession(user string, address string) *ssh.Session {
  client := state.newClient(user, address)

  session, err := client.NewSession()
  if err != nil {
      panic("Failed to create session: " + err.Error())
  }

  return session
}

func (state *AppState) newClient(user string, address string) *ssh.Client {
  config := state.newClientConfig(user)

  client, err := ssh.Dial("tcp", address + ":22", config)
  if err != nil {
      panic("Failed to dial: " + err.Error())
  }

  return client
}

func (state *AppState) newClientConfig(user string) *ssh.ClientConfig {
  return &ssh.ClientConfig{
      User: user,
      Auth: state.Configuration.AuthMethods,
  }
}<commit_msg>Fixed code style<commit_after>package main

import (
	"code.google.com/p/go.crypto/ssh"
	"encoding/json"
	"github.com/go-martini/martini"
	"github.com/martini-contrib/render"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"time"
)

type EndpointConfiguration struct {
	User    string
	Host    string
	Command string
}

type AppConfiguration struct {
	PrivateKeyFile string
	AuthMethods    []ssh.AuthMethod
	Source         EndpointConfiguration
	Destination    EndpointConfiguration
}

type AppState struct {
	Working       bool
	StartTime     *time.Time
	Configuration AppConfiguration
	Logger        *log.Logger
}

func NewAppState() *AppState {
	file, err := os.Open("config.json")
	if err != nil {
		panic("Could not open config.json: " + err.Error())
	}

	decoder := json.NewDecoder(file)

	config := AppConfiguration{}
	if err := decoder.Decode(&config); err != nil {
		panic("Could not decode config.json: " + err.Error())
	}

	privKeyText, err := ioutil.ReadFile(config.PrivateKeyFile)
	if err != nil {
		panic("Failed to read private key: " + err.Error())
	}

	privKey, err := ssh.ParseRawPrivateKey(privKeyText)
	if err != nil {
		panic("Failed to parse private key: " + err.Error())
	}

	signer, err := ssh.NewSignerFromKey(privKey)
	if err != nil {
		panic("Failed to create signer: " + err.Error())
	}

	config.AuthMethods = []ssh.AuthMethod{
		ssh.PublicKeys(signer),
	}

	logger := log.New(os.Stdout, "[big-red] ", 0)

	return &AppState{false, nil, config, logger}
}

func main() {
	state := NewAppState()

	m := martini.Classic()
	m.Use(render.Renderer())

	m.Get("/", func(r render.Render) {
		r.HTML(200, "index", state)
	})

	m.Get("/press", func(r render.Render) {
		if !state.Working {
			go state.PerformDump()
		}

		r.Redirect("/")
	})

	m.Run()
}

func (state *AppState) PerformDump() {
	defer func() {
		if r := recover(); r != nil {
			buf := make([]byte, 1<<16)
			runtime.Stack(buf, false)
			state.Logger.Println(r)
			state.Logger.Println(buf)
		}

		state.Logger.Println("Done. Took " + state.Elapsed().String())

		state.Working = false
		state.StartTime = nil
	}()

	start := time.Now()
	state.Working = true
	state.StartTime = &start

	state.Logger.Println("Started performing work")

	sourceSession := state.newSourceSession()
	defer sourceSession.Close()

	destSession := state.newDestinationSession()
	defer destSession.Close()

	stdoutPipe, err := sourceSession.StdoutPipe()
	if err != nil {
		panic("Could create pipe: " + err.Error())
	}

	destSession.Stdin = stdoutPipe
	destSession.Stdout = os.Stdout
	destSession.Stderr = os.Stderr

	sourceSession.Stderr = os.Stderr

	if err := destSession.Start(state.Configuration.Destination.Command); err != nil {
		panic("Failed to run destination command: " + err.Error())
	}

	if err := sourceSession.Start(state.Configuration.Source.Command); err != nil {
		panic("Failed to run source command: " + err.Error())
	}

	sourceSession.Wait()
	destSession.Wait()
}

func (state *AppState) Elapsed() time.Duration {
	return time.Since(*state.StartTime)
}

func (state *AppState) newSourceSession() *ssh.Session {
	return state.newSession(state.Configuration.Source.User, state.Configuration.Source.Host)
}

func (state *AppState) newDestinationSession() *ssh.Session {
	return state.newSession(state.Configuration.Destination.User, state.Configuration.Destination.Host)
}

func (state *AppState) newSession(user string, address string) *ssh.Session {
	client := state.newClient(user, address)

	session, err := client.NewSession()
	if err != nil {
		panic("Failed to create session: " + err.Error())
	}

	return session
}

func (state *AppState) newClient(user string, address string) *ssh.Client {
	config := state.newClientConfig(user)

	client, err := ssh.Dial("tcp", address+":22", config)
	if err != nil {
		panic("Failed to dial: " + err.Error())
	}

	return client
}

func (state *AppState) newClientConfig(user string) *ssh.ClientConfig {
	return &ssh.ClientConfig{
		User: user,
		Auth: state.Configuration.AuthMethods,
	}
}
<|endoftext|>