<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"image"
	"image/color"
	"image/png"
	"log"
	"os"
	"os/signal"
	"strconv"
	"time"
)

// Version number of this tool
const Version = "0.0.2"

// Struct for each network packet
type Data struct {
	toa     int64 // Timestamp in microseconds
	payload []byte
}

func getBitsFromPacket(packet []byte, byteP, bitP *int, bpP int) uint8 {
	var c uint8
	for i := 0; i < (bpP / 3); i++ {
		if *byteP >= len(packet) {
			break
		}
		c |= (packet[*byteP] & (1 << uint8(7-*bitP)))
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	}
	return c
}

func createPixel(packet []byte, byteP, bitP *int, bpP int) (c color.Color) {
	var r, g, b uint8

	if bpP == 1 {
		if (packet[*byteP] & (1 << uint8(7-*bitP))) == 0 {
			c = color.NRGBA{R: 0,
				G: 0,
				B: 0,
				A: 255}
		} else {
			c = color.NRGBA{R: 255,
				G: 255,
				B: 255,
				A: 255}
		}
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	} else {
		r = getBitsFromPacket(packet, byteP, bitP, bpP)
		g = getBitsFromPacket(packet, byteP, bitP, bpP)
		b = getBitsFromPacket(packet, byteP, bitP, bpP)

		c = color.NRGBA{R: r,
			G: g,
			B: b,
			A: 255}
	}
	return
}

func createTerminalVisualization(data []Data, bitsPerPixel int) {
	var bitPos int
	var bytePos int
	var packetLen int

	for i := range data {
		packetLen = len(data[i].payload)
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[i].payload, &bytePos, &bitPos, bitsPerPixel)
			r, g, b, _ := c.RGBA()
			fmt.Printf("\x1B[0m\x1B[38;2;%d;%d;%dm\u2588", uint8(r), uint8(g), uint8(b))
			if bytePos >= packetLen {
				break
			}
		}
		fmt.Printf("\x1B[m\n")

	}

}
func createTimeVisualization(data []Data, xMax int, prefix string, ts uint, bitsPerPixel int) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int
	var firstPkg time.Time

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/bitsPerPixel+1, int(ts)))

	for pkg := range data {
		if firstPkg.IsZero() {
			firstPkg = time.Unix(0, data[pkg].toa*int64(time.Microsecond))
		}
		packetLen = len(data[pkg].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[pkg].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, int(data[pkg].toa%int64(ts)), c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}
	}

	filename := prefix
	filename += "-"
	filename += firstPkg.Format(time.RFC3339Nano)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	return
}

func createFixedVisualization(data []Data, xMax int, prefix string, num int, bitsPerPixel int) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/bitsPerPixel+1, len(data)))

	for yPos := range data {
		packetLen = len(data[yPos].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[yPos].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, yPos, c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}

	}

	filename := prefix
	filename += strconv.Itoa(num)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	return
}

func handlePackets(ps *gopacket.PacketSource, num uint, ch chan Data, sig <-chan os.Signal) {
	var count uint
	for packet := range ps.Packets() {
		var k Data

		select {
		case isr := <-sig:
			fmt.Println(isr)
			close(ch)
			return
		default:
		}
		count++
		if num != 0 && count > num {
			break
		}

		elements := packet.Data()
		if len(elements) == 0 {
			continue
		}
		k = Data{toa: (packet.Metadata().CaptureInfo.Timestamp.UnixNano() / int64(time.Microsecond)), payload: packet.Data()}
		ch <- k
	}
	close(ch)
	return
}

func availableInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	for _, device := range devices {
		if len(device.Addresses) == 0 {
			continue
		}
		fmt.Println("Interface: ", device.Name)
		for _, address := range device.Addresses {
			fmt.Println("   IP address:  ", address.IP)
			fmt.Println("   Subnet mask: ", address.Netmask)
		}
		fmt.Println("")
	}
}

func main() {
	var err error
	var handle *pcap.Handle
	var data []Data
	var xMax int
	var index int = 1
	var slicer int64
	var flagTimeslize bool = false
	ch := make(chan Data)
	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt)

	dev := flag.String("interface", "", "Choose an interface for online processing")
	file := flag.String("file", "", "Choose a file for offline processing")
	filter := flag.String("filter", "", "Set a specific filter")
	lst := flag.Bool("list_interfaces", false, "List available interfaces")
	vers := flag.Bool("version", false, "Show version")
	help := flag.Bool("help", false, "Show this help")
	terminalOut := flag.Bool("terminal", false, "Visualize on terminal")
	num := flag.Uint("count", 25, "Number of packets to process.\n\tIf argument is 0 the limit is removed")
	output := flag.String("prefix", "image", "Prefix of the resulting image")
	size := flag.Uint("size", 25, "Number of packets per image")
	bits := flag.Uint("bits", 24, "Number of bits per pixel.\n\tIt must be divisible by three and smaller than 25\n\tTo get black/white results, choose 1 as input.")
	ts := flag.Uint("timeslize", 0, "Number of microseconds per resulting image.\n\tSo each pixel of the height of the resulting image represents one microsecond")
	flag.Parse()

	if flag.NFlag() < 1 {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	if *lst {
		availableInterfaces()
		return
	}

	if *vers {
		fmt.Println("Version:", Version)
		return
	}

	if *bits%3 != 0 && *bits != 1 {
		fmt.Println(*bits, "must be divisible by three or one")
		return
	} else if *bits > 25 {
		fmt.Println(*bits, "must be smaller than 25")
		return
	}

	if *ts != 0 {
		flagTimeslize = true
	}

	if *help {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	switch {
	case flagTimeslize == true && *terminalOut:
		fmt.Println("-timeslize and -terminal can't be combined")

		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	if len(*dev) > 0 {
		handle, err = pcap.OpenLive(*dev, 4096, true, pcap.BlockForever)
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	} else if len(*file) > 0 {
		handle, err = pcap.OpenOffline(*file)
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	} else {
		fmt.Println("Source is missing")
		return
	}
	defer handle.Close()

	if len(*filter) != 0 {
		err = handle.SetBPFFilter(*filter)
		if err != nil {
			log.Fatal(err, "\tInvalid filter: ", *filter)
			os.Exit(1)
		}
	}

	packetSource := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
	packetSource.DecodeOptions = gopacket.Lazy

	go handlePackets(packetSource, *num, ch, sig)

	switch {
	case flagTimeslize:
		for i, ok := <-ch; ok; i, ok = <-ch {
			if slicer == 0 {
				slicer = i.toa + int64(*ts)
			}
			if slicer < i.toa {
				xMax++
				createTimeVisualization(data, xMax, *output, *ts, int(*bits))
				xMax = 0
				data = data[:0]
				slicer = i.toa + int64(*ts)
			}
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
		}
	case *terminalOut:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			createTerminalVisualization(data, int(*bits))
			data = data[:0]
		}
	default:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
			if len(data) >= int(*size) {
				xMax++
				createFixedVisualization(data, xMax, *output, index, int(*bits))
				xMax = 0
				index++
				data = data[:0]
			}
		}
	}

	if len(data) > 0 {
		xMax++
		if flagTimeslize {
			createTimeVisualization(data, xMax, *output, *ts, int(*bits))
		} else {
			createFixedVisualization(data, xMax, *output, index, int(*bits))
		}
	}

}
<commit_msg>Add some comments<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"image"
	"image/color"
	"image/png"
	"log"
	"os"
	"os/signal"
	"strconv"
	"time"
)

// Version number of this tool
const Version = "0.0.2"

// Data is a struct for each network packet
type Data struct {
	toa     int64 // Timestamp of arrival in microseconds
	payload []byte  // Copied network packet
}

func getBitsFromPacket(packet []byte, byteP, bitP *int, bpP int) uint8 {
	var c uint8
	for i := 0; i < (bpP / 3); i++ {
		if *byteP >= len(packet) {
			break
		}
		c |= (packet[*byteP] & (1 << uint8(7-*bitP)))
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	}
	return c
}

func createPixel(packet []byte, byteP, bitP *int, bpP int) (c color.Color) {
	var r, g, b uint8

	if bpP == 1 {
		if (packet[*byteP] & (1 << uint8(7-*bitP))) == 0 {
			c = color.NRGBA{R: 0,
				G: 0,
				B: 0,
				A: 255}
		} else {
			c = color.NRGBA{R: 255,
				G: 255,
				B: 255,
				A: 255}
		}
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	} else {
		r = getBitsFromPacket(packet, byteP, bitP, bpP)
		g = getBitsFromPacket(packet, byteP, bitP, bpP)
		b = getBitsFromPacket(packet, byteP, bitP, bpP)

		c = color.NRGBA{R: r,
			G: g,
			B: b,
			A: 255}
	}
	return
}

func createTerminalVisualization(data []Data, bitsPerPixel int) {
	var bitPos int
	var bytePos int
	var packetLen int

	for i := range data {
		packetLen = len(data[i].payload)
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[i].payload, &bytePos, &bitPos, bitsPerPixel)
			r, g, b, _ := c.RGBA()
			fmt.Printf("\x1B[0m\x1B[38;2;%d;%d;%dm\u2588", uint8(r), uint8(g), uint8(b))
			if bytePos >= packetLen {
				break
			}
		}
		fmt.Printf("\x1B[m\n")

	}

}
func createTimeVisualization(data []Data, xMax int, prefix string, ts uint, bitsPerPixel int) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int
	var firstPkg time.Time

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/bitsPerPixel+1, int(ts)))

	for pkg := range data {
		if firstPkg.IsZero() {
			firstPkg = time.Unix(0, data[pkg].toa*int64(time.Microsecond))
		}
		packetLen = len(data[pkg].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[pkg].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, int(data[pkg].toa%int64(ts)), c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}
	}

	filename := prefix
	filename += "-"
	filename += firstPkg.Format(time.RFC3339Nano)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	return
}

func createFixedVisualization(data []Data, xMax int, prefix string, num int, bitsPerPixel int) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/bitsPerPixel+1, len(data)))

	for yPos := range data {
		packetLen = len(data[yPos].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[yPos].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, yPos, c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}

	}

	filename := prefix
	filename += strconv.Itoa(num)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	return
}

func handlePackets(ps *gopacket.PacketSource, num uint, ch chan Data, sig <-chan os.Signal) {
	var count uint
	for packet := range ps.Packets() {
		var k Data

		select {
		case isr := <-sig:
			fmt.Println(isr)
			close(ch)
			return
		default:
		}
		count++
		if num != 0 && count > num {
			break
		}

		elements := packet.Data()
		if len(elements) == 0 {
			continue
		}
		k = Data{toa: (packet.Metadata().CaptureInfo.Timestamp.UnixNano() / int64(time.Microsecond)), payload: packet.Data()}
		ch <- k
	}
	close(ch)
	return
}

func availableInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	for _, device := range devices {
		if len(device.Addresses) == 0 {
			continue
		}
		fmt.Println("Interface: ", device.Name)
		for _, address := range device.Addresses {
			fmt.Println("   IP address:  ", address.IP)
			fmt.Println("   Subnet mask: ", address.Netmask)
		}
		fmt.Println("")
	}
}

func main() {
	var err error
	var handle *pcap.Handle
	var data []Data
	var xMax int
	var index int = 1
	var slicer int64
	var flagTimeslize bool = false
	ch := make(chan Data)
	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt)

	dev := flag.String("interface", "", "Choose an interface for online processing")
	file := flag.String("file", "", "Choose a file for offline processing")
	filter := flag.String("filter", "", "Set a specific filter")
	lst := flag.Bool("list_interfaces", false, "List available interfaces")
	vers := flag.Bool("version", false, "Show version")
	help := flag.Bool("help", false, "Show this help")
	terminalOut := flag.Bool("terminal", false, "Visualize on terminal")
	num := flag.Uint("count", 25, "Number of packets to process.\n\tIf argument is 0 the limit is removed")
	output := flag.String("prefix", "image", "Prefix of the resulting image")
	size := flag.Uint("size", 25, "Number of packets per image")
	bits := flag.Uint("bits", 24, "Number of bits per pixel.\n\tIt must be divisible by three and smaller than 25\n\tTo get black/white results, choose 1 as input.")
	ts := flag.Uint("timeslize", 0, "Number of microseconds per resulting image.\n\tSo each pixel of the height of the resulting image represents one microsecond")
	flag.Parse()

	if flag.NFlag() < 1 {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	if *lst {
		availableInterfaces()
		return
	}

	if *vers {
		fmt.Println("Version:", Version)
		return
	}

	if *bits%3 != 0 && *bits != 1 {
		fmt.Println(*bits, "must be divisible by three or one")
		return
	} else if *bits > 25 {
		fmt.Println(*bits, "must be smaller than 25")
		return
	}

	if *ts != 0 {
		flagTimeslize = true
	}

	if *help {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	switch {
	case flagTimeslize == true && *terminalOut:
		fmt.Println("-timeslize and -terminal can't be combined")

		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	if len(*dev) > 0 {
		handle, err = pcap.OpenLive(*dev, 4096, true, pcap.BlockForever)
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	} else if len(*file) > 0 {
		handle, err = pcap.OpenOffline(*file)
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	} else {
		fmt.Println("Source is missing")
		return
	}
	defer handle.Close()

	if len(*filter) != 0 {
		err = handle.SetBPFFilter(*filter)
		if err != nil {
			log.Fatal(err, "\tInvalid filter: ", *filter)
			os.Exit(1)
		}
	}

	packetSource := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
	packetSource.DecodeOptions = gopacket.Lazy

	go handlePackets(packetSource, *num, ch, sig)

	switch {
	case flagTimeslize:
		for i, ok := <-ch; ok; i, ok = <-ch {
			if slicer == 0 {
				slicer = i.toa + int64(*ts)
			}
			if slicer < i.toa {
				xMax++
				createTimeVisualization(data, xMax, *output, *ts, int(*bits))
				xMax = 0
				data = data[:0]
				slicer = i.toa + int64(*ts)
			}
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
		}
	case *terminalOut:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			createTerminalVisualization(data, int(*bits))
			data = data[:0]
		}
	default:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
			if len(data) >= int(*size) {
				xMax++
				createFixedVisualization(data, xMax, *output, index, int(*bits))
				xMax = 0
				index++
				data = data[:0]
			}
		}
	}

	if len(data) > 0 {
		xMax++
		if flagTimeslize {
			createTimeVisualization(data, xMax, *output, *ts, int(*bits))
		} else {
			createFixedVisualization(data, xMax, *output, index, int(*bits))
		}
	}

}
<|endoftext|>


<commit_before>
package goflickr

import (
	"bytes"
	"crypto/md5"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"sort"
)

const (
	endpoint        = "https://api.flickr.com/services/rest/?"
	uploadEndpoint  = "https://api.flickr.com/services/upload/"
	replaceEndpoint = "https://api.flickr.com/services/replace/"
	authEndpoint    = "https://www.flickr.com/services/auth/?"
	apiHost         = "api.flickr.com"
)

type Request struct {
	ApiKey string
	ApiSecret string
	Method string
	Signature string
	Args   map[string]string
}

type Response struct {
	Status  string         `xml:"stat,attr"`
	Error   *ResponseError `xml:"err"`
	Payload string         `xml:",innerxml"`
}

type ResponseError struct {
	Code    string `xml:"code,attr"`
	Message string `xml:"msg,attr"`
}

type Frob struct {
	Payload string `xml:"frob"`
}

type AuthUser struct {
	Fullname string `xml:"fullname,attr"`
	Nsid string `xml:"nsid,attr"`
	Username string `xml:"username,attr"`
}
type Auth struct {
	Token string `xml:"auth>token"`
	User AuthUser `xml:"auth>user"`
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

type Error string

func (e Error) Error() string {
	return string(e)
}

func (request *Request) GetFrob() Frob {
	request.Method = "flickr.auth.getFrob"
	s, _ := request.doGet(endpoint)
	var f Frob
	xml.Unmarshal(s, &f)
	return f
}

func (request *Request) GetToken(frob Frob) Auth {
	request.Method = "flickr.auth.getToken"
	request.Args = request.getArgsPlusN(1)
	request.Args["frob"] = frob.Payload
	s, _ := request.doGet(endpoint)
	var a Auth
	xml.Unmarshal(s, &a)
	return a
}

func (request *Request) GetSig() string {
	args := request.getArgsPlusN(2)
	args["api_key"] = request.ApiKey
	if request.Method != "" {
		args["method"] = request.Method
	}

	// Sort array keys
	// fixme: got to bet a better way to sort these.
	sorted_keys := make([]string, len(args))
	i := 0
	for k := range args {
		sorted_keys[i] = k
		i++
	}
	sort.Strings(sorted_keys)
	// Build out ordered key-value string prefixed by secret
	s := request.ApiSecret
	for _, key := range sorted_keys {
		if args[key] != "" {
			s += fmt.Sprintf("%s%s", key, args[key])
		}
	}
	fmt.Println(s)
	// Have the full string, now hash
	hash := md5.New()
	hash.Write([]byte(s))
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func (request *Request) Sign(secret string) {
	args := request.Args

	// Remove api_sig
	delete(args, "api_sig")

	sorted_keys := make([]string, len(args)+2)

	args["api_key"] = request.ApiKey
	args["method"] = request.Method

	// Sort array keys
	i := 0
	for k := range args {
		sorted_keys[i] = k
		i++
	}
	sort.Strings(sorted_keys)

	// Build out ordered key-value string prefixed by secret
	s := secret
	for _, key := range sorted_keys {
		if args[key] != "" {
			s += fmt.Sprintf("%s%s", key, args[key])
		}
	}

	// Since we're only adding two keys, it's easier
	// and more space-efficient to just delete them
	// them copy the whole map
	delete(args, "api_key")
	delete(args, "method")

	// Have the full string, now hash
	hash := md5.New()
	hash.Write([]byte(s))

	// Add api_sig as one of the args
	args["api_sig"] = fmt.Sprintf("%x", hash.Sum(nil))
}

// func (request *Request) URL() string {
// 	args := request.Args
// 	args["api_key"] = request.ApiKey
// 	if request.Method != "" {
// 		args["method"] = request.Method
// 	}
// 	if request.Signature != "" {
// 		args["api_sig"] = request.Signature
// 	}
// 	return endpoint + encodeQuery(args)
// }

func (request *Request) getURL(url_base string) string {
	args := request.getArgsPlusN(3)

	args["api_key"] = request.ApiKey
	if request.Method != "" {
		args["method"] = request.Method
	}
	if request.Signature != "" {
		args["api_sig"] = request.Signature
	}
	return url_base + encodeQuery(args)
}

func (request *Request) doGet(earl string) (response []byte, ret error) {
	if request.ApiKey == "" || request.Method == "" {
		return []byte(nil), Error("Need both API key and method")
	}

	request.Signature = request.GetSig()

	s := request.getURL(earl)

	res, err := http.Get(s)
	defer res.Body.Close()
	if err != nil {
		return []byte(nil), err
	}
	request.Signature = ""
	request.Method = ""
	body, _ := ioutil.ReadAll(res.Body)
	return body, nil
}

// func (request *Request) Execute() (response []byte, ret error) {
// 	if request.ApiKey == "" || request.Method == "" {
// 		return []byte(nil), Error("Need both API key and method")
// 	}

// 	request.Signature = request.GetSig()

// 	s := request.URL()

// 	res, err := http.Get(s)
// 	defer res.Body.Close()
// 	if err != nil {
// 		return []byte(nil), err
// 	}

// 	body, _ := ioutil.ReadAll(res.Body)
// 	return body, nil
// }

func encodeQuery(args map[string]string) string {
	i := 0
	s := bytes.NewBuffer(nil)
	for k, v := range args {
		if i != 0 {
			s.WriteString("&")
		}
		i++
		s.WriteString(k + "=" + url.QueryEscape(v))
	}
	return s.String()
}

func (request *Request) buildPost(url_ string, filename string, filetype string) (*http.Request, error) {
	real_url, _ := url.Parse(url_)

	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}

	stat, err := f.Stat()
	if err != nil {
		return nil, err
	}
	f_size := stat.Size()

	request.Args["api_key"] = request.ApiKey

	boundary, end := "----###---###--flickr-go-rules", "\r\n"

	// Build out all of POST body sans file
	header := bytes.NewBuffer(nil)
	for k, v := range request.Args {
		header.WriteString("--" + boundary + end)
		header.WriteString("Content-Disposition: form-data; name=\"" + k + "\"" + end + end)
		header.WriteString(v + end)
	}
	header.WriteString("--" + boundary + end)
	header.WriteString("Content-Disposition: form-data; name=\"photo\"; filename=\"photo.jpg\"" + end)
	header.WriteString("Content-Type: " + filetype + end + end)

	footer := bytes.NewBufferString(end + "--" + boundary + "--" + end)

	body_len := int64(header.Len()) + int64(footer.Len()) + f_size

	r, w := io.Pipe()
	go func() {
		pieces := []io.Reader{header, f, footer}

		for _, k := range pieces {
			_, err = io.Copy(w, k)
			if err != nil {
				w.CloseWithError(nil)
				return
			}
		}
		f.Close()
		w.Close()
	}()

	http_header := make(http.Header)
	http_header.Add("Content-Type", "multipart/form-data; boundary="+boundary)

	postRequest := &http.Request{
		Method:        "POST",
		URL:           real_url,
		Host:          apiHost,
		Header:        http_header,
		Body:          r,
		ContentLength: body_len,
	}
	return postRequest, nil
}

// Example:
// r.Upload("thumb.jpg", "image/jpeg")
func (request *Request) Upload(filename string, filetype string, token string) (response *Response, err error) {
	request.Args = request.getArgsPlusN(5)

	request.Args["is_public"] = "0"
	request.Args["is_family"] = "0"
	request.Args["is_friend"] = "0"
	request.Args["auth_token"] = token
	request.Args["api_sig"] = request.GetSig()
	postRequest, err := request.buildPost(uploadEndpoint, filename, filetype)
	if err != nil {
		return nil, err
	}
	return sendPost(postRequest)
}

func (request *Request) getArgsPlusN(n int) map[string]string {
	args := make(map[string]string, len(request.Args) + n)
	for k, v := range request.Args {
		args[k] = v
	}
	return args
}

func (request *Request) AuthUrl(frob string, perms string) (url string) {
	request.Args = request.getArgsPlusN(2)
	request.Args["frob"] = frob
	request.Args["perms"] = perms
	return request.getURL(authEndpoint)  + "&api_sig=" + request.GetSig()
}

func (request *Request) Replace(filename string, filetype string) (response *Response, err error) {
	postRequest, err := request.buildPost(replaceEndpoint, filename, filetype)
	if err != nil {
		return nil, err
	}
	return sendPost(postRequest)
}

func sendPost(postRequest *http.Request) (response *Response, err error) {
	// Create and use TCP connection (lifted mostly wholesale from http.send)
	client := &http.DefaultClient
	resp, err := client.Do(postRequest)

	if err != nil {
		return nil, err
	}
	rawBody, _ := ioutil.ReadAll(resp.Body)
	resp.Body.Close()

	var r Response
	err = xml.Unmarshal(rawBody, &r)

	return &r, err
}

<commit_msg>deleted some stuff I'm not using anymore.<commit_after>
package goflickr

import (
	"bytes"
	"crypto/md5"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"sort"
)

const (
	endpoint        = "https://api.flickr.com/services/rest/?"
	uploadEndpoint  = "https://api.flickr.com/services/upload/"
	replaceEndpoint = "https://api.flickr.com/services/replace/"
	authEndpoint    = "https://www.flickr.com/services/auth/?"
	apiHost         = "api.flickr.com"
)

type Request struct {
	ApiKey string
	ApiSecret string
	Method string
	Signature string
	Args   map[string]string
}

type Response struct {
	Status  string         `xml:"stat,attr"`
	Error   *ResponseError `xml:"err"`
	Payload string         `xml:",innerxml"`
}

type ResponseError struct {
	Code    string `xml:"code,attr"`
	Message string `xml:"msg,attr"`
}

type Frob struct {
	Payload string `xml:"frob"`
}

type AuthUser struct {
	Fullname string `xml:"fullname,attr"`
	Nsid string `xml:"nsid,attr"`
	Username string `xml:"username,attr"`
}
type Auth struct {
	Token string `xml:"auth>token"`
	User AuthUser `xml:"auth>user"`
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

type Error string

func (e Error) Error() string {
	return string(e)
}

func (request *Request) GetFrob() Frob {
	request.Method = "flickr.auth.getFrob"
	s, _ := request.doGet(endpoint)
	var f Frob
	xml.Unmarshal(s, &f)
	return f
}

func (request *Request) GetToken(frob Frob) Auth {
	request.Method = "flickr.auth.getToken"
	request.Args = request.getArgsPlusN(1)
	request.Args["frob"] = frob.Payload
	s, _ := request.doGet(endpoint)
	var a Auth
	xml.Unmarshal(s, &a)
	return a
}

func (request *Request) GetSig() string {
	args := request.getArgsPlusN(2)
	args["api_key"] = request.ApiKey
	if request.Method != "" {
		args["method"] = request.Method
	}

	// Sort array keys
	// fixme: got to bet a better way to sort these.
	sorted_keys := make([]string, len(args))
	i := 0
	for k := range args {
		sorted_keys[i] = k
		i++
	}
	sort.Strings(sorted_keys)
	// Build out ordered key-value string prefixed by secret
	s := request.ApiSecret
	for _, key := range sorted_keys {
		if args[key] != "" {
			s += fmt.Sprintf("%s%s", key, args[key])
		}
	}
	fmt.Println(s)
	// Have the full string, now hash
	hash := md5.New()
	hash.Write([]byte(s))
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func (request *Request) getURL(url_base string) string {
	args := request.getArgsPlusN(3)

	args["api_key"] = request.ApiKey
	if request.Method != "" {
		args["method"] = request.Method
	}
	if request.Signature != "" {
		args["api_sig"] = request.Signature
	}
	return url_base + encodeQuery(args)
}

func (request *Request) doGet(earl string) (response []byte, ret error) {
	if request.ApiKey == "" || request.Method == "" {
		return []byte(nil), Error("Need both API key and method")
	}

	request.Signature = request.GetSig()

	s := request.getURL(earl)

	res, err := http.Get(s)
	defer res.Body.Close()
	if err != nil {
		return []byte(nil), err
	}
	request.Signature = ""
	request.Method = ""
	body, _ := ioutil.ReadAll(res.Body)
	return body, nil
}

func encodeQuery(args map[string]string) string {
	i := 0
	s := bytes.NewBuffer(nil)
	for k, v := range args {
		if i != 0 {
			s.WriteString("&")
		}
		i++
		s.WriteString(k + "=" + url.QueryEscape(v))
	}
	return s.String()
}

func (request *Request) buildPost(url_ string, filename string, filetype string) (*http.Request, error) {
	real_url, _ := url.Parse(url_)

	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}

	stat, err := f.Stat()
	if err != nil {
		return nil, err
	}
	f_size := stat.Size()

	request.Args["api_key"] = request.ApiKey

	boundary, end := "----###---###--flickr-go-rules", "\r\n"

	// Build out all of POST body sans file
	header := bytes.NewBuffer(nil)
	for k, v := range request.Args {
		header.WriteString("--" + boundary + end)
		header.WriteString("Content-Disposition: form-data; name=\"" + k + "\"" + end + end)
		header.WriteString(v + end)
	}
	header.WriteString("--" + boundary + end)
	header.WriteString("Content-Disposition: form-data; name=\"photo\"; filename=\"photo.jpg\"" + end)
	header.WriteString("Content-Type: " + filetype + end + end)

	footer := bytes.NewBufferString(end + "--" + boundary + "--" + end)

	body_len := int64(header.Len()) + int64(footer.Len()) + f_size

	r, w := io.Pipe()
	go func() {
		pieces := []io.Reader{header, f, footer}

		for _, k := range pieces {
			_, err = io.Copy(w, k)
			if err != nil {
				w.CloseWithError(nil)
				return
			}
		}
		f.Close()
		w.Close()
	}()

	http_header := make(http.Header)
	http_header.Add("Content-Type", "multipart/form-data; boundary="+boundary)

	postRequest := &http.Request{
		Method:        "POST",
		URL:           real_url,
		Host:          apiHost,
		Header:        http_header,
		Body:          r,
		ContentLength: body_len,
	}
	return postRequest, nil
}

// Example:
// r.Upload("thumb.jpg", "image/jpeg")
func (request *Request) Upload(filename string, filetype string, token string) (response *Response, err error) {
	request.Args = request.getArgsPlusN(5)

	request.Args["is_public"] = "0"
	request.Args["is_family"] = "0"
	request.Args["is_friend"] = "0"
	request.Args["auth_token"] = token
	request.Args["api_sig"] = request.GetSig()
	postRequest, err := request.buildPost(uploadEndpoint, filename, filetype)
	if err != nil {
		return nil, err
	}
	return sendPost(postRequest)
}

func (request *Request) getArgsPlusN(n int) map[string]string {
	args := make(map[string]string, len(request.Args) + n)
	for k, v := range request.Args {
		args[k] = v
	}
	return args
}

func (request *Request) AuthUrl(frob string, perms string) (url string) {
	request.Args = request.getArgsPlusN(2)
	request.Args["frob"] = frob
	request.Args["perms"] = perms
	return request.getURL(authEndpoint)  + "&api_sig=" + request.GetSig()
}

func (request *Request) Replace(filename string, filetype string) (response *Response, err error) {
	postRequest, err := request.buildPost(replaceEndpoint, filename, filetype)
	if err != nil {
		return nil, err
	}
	return sendPost(postRequest)
}

func sendPost(postRequest *http.Request) (response *Response, err error) {
	// Create and use TCP connection (lifted mostly wholesale from http.send)
	client := &http.DefaultClient
	resp, err := client.Do(postRequest)

	if err != nil {
		return nil, err
	}
	rawBody, _ := ioutil.ReadAll(resp.Body)
	resp.Body.Close()

	var r Response
	err = xml.Unmarshal(rawBody, &r)

	return &r, err
}

<|endoftext|>


<commit_before>package main

import (
	"./ganglia/api"
	"./ganglia/response"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	// filters
	node    = kingpin.Flag("node", "Node name to search").String()
	cluster = kingpin.Flag("cluster", "Cluster name to search").String()
	metric  = kingpin.Flag("metric", "Metric name to search").String()
	// conditions
	threshhold = kingpin.Flag("threshhold", "Threshhold value").Int()
	condition  = kingpin.Flag("condition", "Condition").Enum("eq", "gt", "ge", "lt", "le")
	// source of the xml
	localPath  = kingpin.Flag("file", "Read gmeta response from local file").ExistingFile()
	remoteIP   = kingpin.Flag("host", "Read gmeta response from host, default: 127.0.0.1").Default("127.0.0.1").String()
	remotePort = kingpin.Flag("port", "Read gmeta response from port, default: 8651").Default("8651").Int()
	// list objects
	listMetric   = kingpin.Flag("list-metrics", "List metrics").Bool()
	listNodes    = kingpin.Flag("list-nodes", "List nodes").Bool()
	listClusters = kingpin.Flag("list-clusters", "List clusters").Bool()
)

func main() {

	kingpin.Parse()

	gResp := response.GMetaResponse{}
	filter := response.MetricFilter{*metric, *node, *cluster, *condition, *threshhold}

	if *localPath != "" {
		gResp = api.ParseFile(*localPath)
	} else {
		gResp = api.ParseSocket(*remoteIP, *remotePort)
	}

	s := api.NewGSet()
	for _, metric := range gResp.Find(&filter) {
		if *listClusters {
			s.PrintIfNotInSet(metric.GCluster.Name)
		}
		if *listNodes {
			s.PrintIfNotInSet(metric.GHost.Name)
		}
		if *listMetric {
			s.PrintIfNotInSet(metric.GMetric.Name)
		}
	}

}
<commit_msg>added deamon mode<commit_after>package main

import (
	"./gogalert/deamon"
	"./gogalert/ganglia/api"
	"./gogalert/ganglia/response"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	// filters
	node    = kingpin.Flag("node", "Node name to search").String()
	cluster = kingpin.Flag("cluster", "Cluster name to search").String()
	metric  = kingpin.Flag("metric", "Metric name to search").String()
	// conditions
	threshhold = kingpin.Flag("threshhold", "Threshhold value").Int()
	condition  = kingpin.Flag("condition", "Condition").Enum("eq", "gt", "ge", "lt", "le")
	// source of the xml
	localPath  = kingpin.Flag("file", "Read gmeta response from local file").ExistingFile()
	remoteHost = kingpin.Flag("host", "Read gmeta response from host, default: 127.0.0.1").Default("127.0.0.1").String()
	remotePort = kingpin.Flag("port", "Read gmeta response from port, default: 8651").Default("8651").Int()
	// list objects
	listMetric   = kingpin.Flag("list-metrics", "List metrics").Bool()
	listNodes    = kingpin.Flag("list-nodes", "List nodes").Bool()
	listClusters = kingpin.Flag("list-clusters", "List clusters").Bool()
	// deamon mode
	deamonMode = kingpin.Flag("deamon", "Run in deamon mode").Bool()
)

func main() {

	kingpin.Parse()

	filter := &response.MetricFilter{*metric, *node, *cluster, *condition, *threshhold}
	source := &api.DataSource{*localPath, *remoteHost, *remotePort}

	if *deamonMode {
		deamon.NewStatServer(source, filter).Start()
	} else {
		s := api.NewGSet()
		for _, metric := range api.Parse(source).Find(filter) {
			if *listClusters {
				s.PrintIfNotInSet(metric.GCluster.Name)
			}
			if *listNodes {
				s.PrintIfNotInSet(metric.GHost.Name)
			}
			if *listMetric {
				s.PrintIfNotInSet(metric.GMetric.Name)
			}
		}
	}
}
<|endoftext|>


<commit_before>package graceful

import (
	"crypto/tls"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"golang.org/x/net/netutil"
)

// Server wraps an http.Server with graceful connection handling.
// It may be used directly in the same way as http.Server, or may
// be constructed with the global functions in this package.
//
// Example:
//	srv := &graceful.Server{
//		Timeout: 5 * time.Second,
//		Server: &http.Server{Addr: ":1234", Handler: handler},
//	}
//	srv.ListenAndServe()
type Server struct {
	*http.Server

	// Timeout is the duration to allow outstanding requests to survive
	// before forcefully terminating them.
	Timeout time.Duration

	// Limit the number of outstanding requests
	ListenLimit int

	// ConnState specifies an optional callback function that is
	// called when a client connection changes state. This is a proxy
	// to the underlying http.Server's ConnState, and the original
	// must not be set directly.
	ConnState func(net.Conn, http.ConnState)

	// BeforeShutdown is an optional callback function that is called
	// before the listener is closed.
	BeforeShutdown func()

	// ShutdownInitiated is an optional callback function that is called
	// when shutdown is initiated. It can be used to notify the client
	// side of long lived connections (e.g. websockets) to reconnect.
	ShutdownInitiated func()

	// NoSignalHandling prevents graceful from automatically shutting down
	// on SIGINT and SIGTERM. If set to true, you must shut down the server
	// manually with Stop().
	NoSignalHandling bool

	// Logger used to notify of errors on startup and on stop.
	Logger *log.Logger

	// Interrupted is true if the server is handling a SIGINT or SIGTERM
	// signal and is thus shutting down.
	Interrupted bool

	// interrupt signals the listener to stop serving connections,
	// and the server to shut down.
	interrupt chan os.Signal

	// stopLock is used to protect against concurrent calls to Stop
	stopLock sync.Mutex

	// stopChan is the channel on which callers may block while waiting for
	// the server to stop.
	stopChan chan struct{}

	// chanLock is used to protect access to the various channel constructors.
	chanLock sync.RWMutex

	// connections holds all connections managed by graceful
	connections map[net.Conn]struct{}
}

// Run serves the http.Handler with graceful shutdown enabled.
//
// timeout is the duration to wait until killing active requests and stopping the server.
// If timeout is 0, the server never times out. It waits for all active requests to finish.
func Run(addr string, timeout time.Duration, n http.Handler) {
	srv := &Server{
		Timeout: timeout,
		Server:  &http.Server{Addr: addr, Handler: n},
		Logger:  DefaultLogger(),
	}

	if err := srv.ListenAndServe(); err != nil {
		if opErr, ok := err.(*net.OpError); !ok || (ok && opErr.Op != "accept") {
			srv.Logger.Fatal(err)
		}
	}

}

// RunWithErr is an alternative version of Run function which can return error.
//
// Unlike Run this version will not exit the program if an error is encountered but will
// return it instead.
func RunWithErr(addr string, timeout time.Duration, n http.Handler) error {
	srv := &Server{
		Timeout: timeout,
		Server:  &http.Server{Addr: addr, Handler: n},
		Logger:  DefaultLogger(),
	}

	return srv.ListenAndServe()
}

// ListenAndServe is equivalent to http.Server.ListenAndServe with graceful shutdown enabled.
//
// timeout is the duration to wait until killing active requests and stopping the server.
// If timeout is 0, the server never times out. It waits for all active requests to finish.
func ListenAndServe(server *http.Server, timeout time.Duration) error {
	srv := &Server{Timeout: timeout, Server: server, Logger: DefaultLogger()}
	return srv.ListenAndServe()
}

// ListenAndServe is equivalent to http.Server.ListenAndServe with graceful shutdown enabled.
func (srv *Server) ListenAndServe() error {
	// Create the listener so we can control their lifetime
	addr := srv.Addr
	if addr == "" {
		addr = ":http"
	}
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	return srv.Serve(l)
}

// ListenAndServeTLS is equivalent to http.Server.ListenAndServeTLS with graceful shutdown enabled.
//
// timeout is the duration to wait until killing active requests and stopping the server.
// If timeout is 0, the server never times out. It waits for all active requests to finish.
func ListenAndServeTLS(server *http.Server, certFile, keyFile string, timeout time.Duration) error {
	srv := &Server{Timeout: timeout, Server: server, Logger: DefaultLogger()}
	return srv.ListenAndServeTLS(certFile, keyFile)
}

// ListenTLS is a convenience method that creates an https listener using the
// provided cert and key files. Use this method if you need access to the
// listener object directly. When ready, pass it to the Serve method.
func (srv *Server) ListenTLS(certFile, keyFile string) (net.Listener, error) {
	// Create the listener ourselves so we can control its lifetime
	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

	config := &tls.Config{}
	if srv.TLSConfig != nil {
		*config = *srv.TLSConfig
	}
	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return nil, err
	}

	conn, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}

	tlsListener := tls.NewListener(conn, config)
	return tlsListener, nil
}

// ListenAndServeTLS is equivalent to http.Server.ListenAndServeTLS with graceful shutdown enabled.
func (srv *Server) ListenAndServeTLS(certFile, keyFile string) error {
	l, err := srv.ListenTLS(certFile, keyFile)
	if err != nil {
		return err
	}

	return srv.Serve(l)
}

// ListenAndServeTLSConfig can be used with an existing TLS config and is equivalent to
// http.Server.ListenAndServeTLS with graceful shutdown enabled,
func (srv *Server) ListenAndServeTLSConfig(config *tls.Config) error {
	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

	conn, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(conn, config)
	return srv.Serve(tlsListener)
}

// Serve is equivalent to http.Server.Serve with graceful shutdown enabled.
//
// timeout is the duration to wait until killing active requests and stopping the server.
// If timeout is 0, the server never times out. It waits for all active requests to finish.
func Serve(server *http.Server, l net.Listener, timeout time.Duration) error {
	srv := &Server{Timeout: timeout, Server: server, Logger: DefaultLogger()}
	return srv.Serve(l)
}

// Serve is equivalent to http.Server.Serve with graceful shutdown enabled.
func (srv *Server) Serve(listener net.Listener) error {

	if srv.ListenLimit != 0 {
		listener = netutil.LimitListener(listener, srv.ListenLimit)
	}

	// Track connection state
	add := make(chan net.Conn)
	remove := make(chan net.Conn)

	srv.Server.ConnState = func(conn net.Conn, state http.ConnState) {
		switch state {
		case http.StateNew:
			add <- conn
		case http.StateClosed, http.StateHijacked:
			remove <- conn
		}
		if srv.ConnState != nil {
			srv.ConnState(conn, state)
		}
	}

	// Manage open connections
	shutdown := make(chan chan struct{})
	kill := make(chan struct{})
	go srv.manageConnections(add, remove, shutdown, kill)

	interrupt := srv.interruptChan()
	// Set up the interrupt handler
	if !srv.NoSignalHandling {
		signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM)
	}
	quitting := make(chan struct{})
	go srv.handleInterrupt(interrupt, quitting, listener)

	// Serve with graceful listener.
	// Execution blocks here until listener.Close() is called, above.
	err := srv.Server.Serve(listener)
	if err != nil {
		// If the underlying listening is closed, Serve returns an error
		// complaining about listening on a closed socket. This is expected, so
		// let's ignore the error if we are the ones who explicitly closed the
		// socket.
		select {
		case <-quitting:
			err = nil
		default:
		}
	}

	srv.shutdown(shutdown, kill)

	return err
}

// Stop instructs the type to halt operations and close
// the stop channel when it is finished.
//
// timeout is grace period for which to wait before shutting
// down the server. The timeout value passed here will override the
// timeout given when constructing the server, as this is an explicit
// command to stop the server.
func (srv *Server) Stop(timeout time.Duration) {
	srv.stopLock.Lock()
	defer srv.stopLock.Unlock()

	srv.Timeout = timeout
	interrupt := srv.interruptChan()
	interrupt <- syscall.SIGINT
}

// StopChan gets the stop channel which will block until
// stopping has completed, at which point it is closed.
// Callers should never close the stop channel.
func (srv *Server) StopChan() <-chan struct{} {
	srv.chanLock.Lock()
	defer srv.chanLock.Unlock()

	if srv.stopChan == nil {
		srv.stopChan = make(chan struct{})
	}
	return srv.stopChan
}

// DefaultLogger returns the logger used by Run, RunWithErr, ListenAndServe, ListenAndServeTLS and Serve.
// The logger outputs to STDERR by default.
func DefaultLogger() *log.Logger {
	return log.New(os.Stderr, "[graceful] ", 0)
}

func (srv *Server) manageConnections(add, remove chan net.Conn, shutdown chan chan struct{}, kill chan struct{}) {
	var done chan struct{}
	srv.connections = map[net.Conn]struct{}{}
	for {
		select {
		case conn := <-add:
			srv.connections[conn] = struct{}{}
		case conn := <-remove:
			delete(srv.connections, conn)
			if done != nil && len(srv.connections) == 0 {
				done <- struct{}{}
				return
			}
		case done = <-shutdown:
			if len(srv.connections) == 0 {
				done <- struct{}{}
				return
			}
		case <-kill:
			for k := range srv.connections {
				if err := k.Close(); err != nil {
					srv.log("[ERROR] %s", err)
				}
			}
			return
		}
	}
}

func (srv *Server) interruptChan() chan os.Signal {
	srv.chanLock.Lock()
	defer srv.chanLock.Unlock()

	if srv.interrupt == nil {
		srv.interrupt = make(chan os.Signal, 1)
	}

	return srv.interrupt
}

func (srv *Server) handleInterrupt(interrupt chan os.Signal, quitting chan struct{}, listener net.Listener) {
	for _ = range interrupt {
		if srv.Interrupted {
			srv.log("already shutting down")
			continue
		}
		srv.log("shutdown initiated")
		srv.Interrupted = true
		if srv.BeforeShutdown != nil {
			srv.BeforeShutdown()
		}

		close(quitting)
		srv.SetKeepAlivesEnabled(false)
		if err := listener.Close(); err != nil {
			srv.log("[ERROR] %s", err)
		}

		if srv.ShutdownInitiated != nil {
			srv.ShutdownInitiated()
		}
	}
}

func (srv *Server) log(fmt string, v ...interface{}) {
	if srv.Logger != nil {
		srv.Logger.Printf(fmt, v...)
	}
}

func (srv *Server) shutdown(shutdown chan chan struct{}, kill chan struct{}) {
	// Request done notification
	done := make(chan struct{})
	shutdown <- done

	if srv.Timeout > 0 {
		select {
		case <-done:
		case <-time.After(srv.Timeout):
			close(kill)
		}
	} else {
		<-done
	}
	// Close the stopChan to wake up any blocked goroutines.
	srv.chanLock.Lock()
	if srv.stopChan != nil {
		close(srv.stopChan)
	}
	srv.chanLock.Unlock()
}
<commit_msg>Add tcpKeepAliveListener to clean up connections where the client has disappeared.<commit_after>package graceful

import (
	"crypto/tls"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"golang.org/x/net/netutil"
)

// Server wraps an http.Server with graceful connection handling.
// It may be used directly in the same way as http.Server, or may
// be constructed with the global functions in this package.
//
// Example:
//	srv := &graceful.Server{
//		Timeout: 5 * time.Second,
//		Server: &http.Server{Addr: ":1234", Handler: handler},
//	}
//	srv.ListenAndServe()
type Server struct {
	*http.Server

	// Timeout is the duration to allow outstanding requests to survive
	// before forcefully terminating them.
	Timeout time.Duration

	// Limit the number of outstanding requests
	ListenLimit int

	// TCPKeepAlive sets the TCP keep-alive timeouts on accepted
	// connections. It prunes dead TCP connections ( e.g. closing
	// laptop mid-download)
	TCPKeepAlive time.Duration

	// ConnState specifies an optional callback function that is
	// called when a client connection changes state. This is a proxy
	// to the underlying http.Server's ConnState, and the original
	// must not be set directly.
	ConnState func(net.Conn, http.ConnState)

	// BeforeShutdown is an optional callback function that is called
	// before the listener is closed.
	BeforeShutdown func()

	// ShutdownInitiated is an optional callback function that is called
	// when shutdown is initiated. It can be used to notify the client
	// side of long lived connections (e.g. websockets) to reconnect.
	ShutdownInitiated func()

	// NoSignalHandling prevents graceful from automatically shutting down
	// on SIGINT and SIGTERM. If set to true, you must shut down the server
	// manually with Stop().
	NoSignalHandling bool

	// Logger used to notify of errors on startup and on stop.
	Logger *log.Logger

	// Interrupted is true if the server is handling a SIGINT or SIGTERM
	// signal and is thus shutting down.
	Interrupted bool

	// interrupt signals the listener to stop serving connections,
	// and the server to shut down.
	interrupt chan os.Signal

	// stopLock is used to protect against concurrent calls to Stop
	stopLock sync.Mutex

	// stopChan is the channel on which callers may block while waiting for
	// the server to stop.
	stopChan chan struct{}

	// chanLock is used to protect access to the various channel constructors.
	chanLock sync.RWMutex

	// connections holds all connections managed by graceful
	connections map[net.Conn]struct{}
}

// Run serves the http.Handler with graceful shutdown enabled.
//
// timeout is the duration to wait until killing active requests and stopping the server.
// If timeout is 0, the server never times out. It waits for all active requests to finish.
func Run(addr string, timeout time.Duration, n http.Handler) {
	srv := &Server{
		Timeout:      timeout,
		TCPKeepAlive: 3 * time.Minute,
		Server:       &http.Server{Addr: addr, Handler: n},
		Logger:       DefaultLogger(),
	}

	if err := srv.ListenAndServe(); err != nil {
		if opErr, ok := err.(*net.OpError); !ok || (ok && opErr.Op != "accept") {
			srv.Logger.Fatal(err)
		}
	}

}

// RunWithErr is an alternative version of Run function which can return error.
//
// Unlike Run this version will not exit the program if an error is encountered but will
// return it instead.
func RunWithErr(addr string, timeout time.Duration, n http.Handler) error {
	srv := &Server{
		Timeout:      timeout,
		TCPKeepAlive: 3 * time.Minute,
		Server:       &http.Server{Addr: addr, Handler: n},
		Logger:       DefaultLogger(),
	}

	return srv.ListenAndServe()
}

// ListenAndServe is equivalent to http.Server.ListenAndServe with graceful shutdown enabled.
//
// timeout is the duration to wait until killing active requests and stopping the server.
// If timeout is 0, the server never times out. It waits for all active requests to finish.
func ListenAndServe(server *http.Server, timeout time.Duration) error {
	srv := &Server{Timeout: timeout, Server: server, Logger: DefaultLogger()}
	return srv.ListenAndServe()
}

// ListenAndServe is equivalent to http.Server.ListenAndServe with graceful shutdown enabled.
func (srv *Server) ListenAndServe() error {
	// Create the listener so we can control their lifetime
	addr := srv.Addr
	if addr == "" {
		addr = ":http"
	}
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	return srv.Serve(l)
}

// ListenAndServeTLS is equivalent to http.Server.ListenAndServeTLS with graceful shutdown enabled.
//
// timeout is the duration to wait until killing active requests and stopping the server.
// If timeout is 0, the server never times out. It waits for all active requests to finish.
func ListenAndServeTLS(server *http.Server, certFile, keyFile string, timeout time.Duration) error {
	srv := &Server{Timeout: timeout, Server: server, Logger: DefaultLogger()}
	return srv.ListenAndServeTLS(certFile, keyFile)
}

// ListenTLS is a convenience method that creates an https listener using the
// provided cert and key files. Use this method if you need access to the
// listener object directly. When ready, pass it to the Serve method.
func (srv *Server) ListenTLS(certFile, keyFile string) (net.Listener, error) {
	// Create the listener ourselves so we can control its lifetime
	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

	config := &tls.Config{}
	if srv.TLSConfig != nil {
		*config = *srv.TLSConfig
	}
	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return nil, err
	}

	conn, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}

	tlsListener := tls.NewListener(conn, config)
	return tlsListener, nil
}

// ListenAndServeTLS is equivalent to http.Server.ListenAndServeTLS with graceful shutdown enabled.
func (srv *Server) ListenAndServeTLS(certFile, keyFile string) error {
	l, err := srv.ListenTLS(certFile, keyFile)
	if err != nil {
		return err
	}

	return srv.Serve(l)
}

// ListenAndServeTLSConfig can be used with an existing TLS config and is equivalent to
// http.Server.ListenAndServeTLS with graceful shutdown enabled,
func (srv *Server) ListenAndServeTLSConfig(config *tls.Config) error {
	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

	conn, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(conn, config)
	return srv.Serve(tlsListener)
}

// Serve is equivalent to http.Server.Serve with graceful shutdown enabled.
//
// timeout is the duration to wait until killing active requests and stopping the server.
// If timeout is 0, the server never times out. It waits for all active requests to finish.
func Serve(server *http.Server, l net.Listener, timeout time.Duration) error {
	srv := &Server{Timeout: timeout, Server: server, Logger: DefaultLogger()}
	return srv.Serve(l)
}

// Serve is equivalent to http.Server.Serve with graceful shutdown enabled.
func (srv *Server) Serve(listener net.Listener) error {

	if srv.ListenLimit != 0 {
		listener = netutil.LimitListener(listener, srv.ListenLimit)
	}

	if srv.TCPKeepAlive != 0 {
		listener = tcpKeepAliveListener{listener.(*net.TCPListener), srv.TCPKeepAlive}
	}

	// Track connection state
	add := make(chan net.Conn)
	remove := make(chan net.Conn)

	srv.Server.ConnState = func(conn net.Conn, state http.ConnState) {
		switch state {
		case http.StateNew:
			add <- conn
		case http.StateClosed, http.StateHijacked:
			remove <- conn
		}
		if srv.ConnState != nil {
			srv.ConnState(conn, state)
		}
	}

	// Manage open connections
	shutdown := make(chan chan struct{})
	kill := make(chan struct{})
	go srv.manageConnections(add, remove, shutdown, kill)

	interrupt := srv.interruptChan()
	// Set up the interrupt handler
	if !srv.NoSignalHandling {
		signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM)
	}
	quitting := make(chan struct{})
	go srv.handleInterrupt(interrupt, quitting, listener)

	// Serve with graceful listener.
	// Execution blocks here until listener.Close() is called, above.
	err := srv.Server.Serve(listener)
	if err != nil {
		// If the underlying listening is closed, Serve returns an error
		// complaining about listening on a closed socket. This is expected, so
		// let's ignore the error if we are the ones who explicitly closed the
		// socket.
		select {
		case <-quitting:
			err = nil
		default:
		}
	}

	srv.shutdown(shutdown, kill)

	return err
}

// Stop instructs the type to halt operations and close
// the stop channel when it is finished.
//
// timeout is grace period for which to wait before shutting
// down the server. The timeout value passed here will override the
// timeout given when constructing the server, as this is an explicit
// command to stop the server.
func (srv *Server) Stop(timeout time.Duration) {
	srv.stopLock.Lock()
	defer srv.stopLock.Unlock()

	srv.Timeout = timeout
	interrupt := srv.interruptChan()
	interrupt <- syscall.SIGINT
}

// StopChan gets the stop channel which will block until
// stopping has completed, at which point it is closed.
// Callers should never close the stop channel.
func (srv *Server) StopChan() <-chan struct{} {
	srv.chanLock.Lock()
	defer srv.chanLock.Unlock()

	if srv.stopChan == nil {
		srv.stopChan = make(chan struct{})
	}
	return srv.stopChan
}

// DefaultLogger returns the logger used by Run, RunWithErr, ListenAndServe, ListenAndServeTLS and Serve.
// The logger outputs to STDERR by default.
func DefaultLogger() *log.Logger {
	return log.New(os.Stderr, "[graceful] ", 0)
}

func (srv *Server) manageConnections(add, remove chan net.Conn, shutdown chan chan struct{}, kill chan struct{}) {
	var done chan struct{}
	srv.connections = map[net.Conn]struct{}{}
	for {
		select {
		case conn := <-add:
			srv.connections[conn] = struct{}{}
		case conn := <-remove:
			delete(srv.connections, conn)
			if done != nil && len(srv.connections) == 0 {
				done <- struct{}{}
				return
			}
		case done = <-shutdown:
			if len(srv.connections) == 0 {
				done <- struct{}{}
				return
			}
		case <-kill:
			for k := range srv.connections {
				if err := k.Close(); err != nil {
					srv.log("[ERROR] %s", err)
				}
			}
			return
		}
	}
}

func (srv *Server) interruptChan() chan os.Signal {
	srv.chanLock.Lock()
	defer srv.chanLock.Unlock()

	if srv.interrupt == nil {
		srv.interrupt = make(chan os.Signal, 1)
	}

	return srv.interrupt
}

func (srv *Server) handleInterrupt(interrupt chan os.Signal, quitting chan struct{}, listener net.Listener) {
	for _ = range interrupt {
		if srv.Interrupted {
			srv.log("already shutting down")
			continue
		}
		srv.log("shutdown initiated")
		srv.Interrupted = true
		if srv.BeforeShutdown != nil {
			srv.BeforeShutdown()
		}

		close(quitting)
		srv.SetKeepAlivesEnabled(false)
		if err := listener.Close(); err != nil {
			srv.log("[ERROR] %s", err)
		}

		if srv.ShutdownInitiated != nil {
			srv.ShutdownInitiated()
		}
	}
}

func (srv *Server) log(fmt string, v ...interface{}) {
	if srv.Logger != nil {
		srv.Logger.Printf(fmt, v...)
	}
}

func (srv *Server) shutdown(shutdown chan chan struct{}, kill chan struct{}) {
	// Request done notification
	done := make(chan struct{})
	shutdown <- done

	if srv.Timeout > 0 {
		select {
		case <-done:
		case <-time.After(srv.Timeout):
			close(kill)
		}
	} else {
		<-done
	}
	// Close the stopChan to wake up any blocked goroutines.
	srv.chanLock.Lock()
	if srv.stopChan != nil {
		close(srv.stopChan)
	}
	srv.chanLock.Unlock()
}

// tcpKeepAliveListener sets TCP keep-alive timeouts on accepted
// connections. It's used by ListenAndServe and ListenAndServeTLS so
// dead TCP connections (e.g. closing laptop mid-download) eventually
// go away.
type tcpKeepAliveListener struct {
	*net.TCPListener
	keepAlivePeriod time.Duration
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(ln.keepAlivePeriod)
	return tc, nil
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package lib

import (
	"github.com/loadimpact/k6/lib/metrics"
	"github.com/loadimpact/k6/stats"
	"net"
	"net/http/httptrace"
	"time"
)

// A Trail represents detailed information about an HTTP request.
// You'd typically get one from a Tracer.
type Trail struct {
	// All metrics will be tagged with this timestamp.
	StartTime time.Time

	// Total request duration, excluding DNS lookup and connect time.
	Duration time.Duration

	Blocked    time.Duration // Waiting to acquire a connection.
	LookingUp  time.Duration // Looking up DNS records.
	Connecting time.Duration // Connecting to remote host.
	Sending    time.Duration // Writing request.
	Waiting    time.Duration // Waiting for first byte.
	Receiving  time.Duration // Receiving response.

	// Detailed connection information.
	ConnReused     bool
	ConnRemoteAddr net.Addr
}

func (tr Trail) Samples(tags map[string]string) []stats.Sample {
	return []stats.Sample{
		{Metric: metrics.HTTPReqs, Time: tr.StartTime, Tags: tags, Value: 1},
		{Metric: metrics.HTTPReqDuration, Time: tr.StartTime, Tags: tags, Value: stats.D(tr.Duration)},
		{Metric: metrics.HTTPReqBlocked, Time: tr.StartTime, Tags: tags, Value: stats.D(tr.Blocked)},
		{Metric: metrics.HTTPReqLookingUp, Time: tr.StartTime, Tags: tags, Value: stats.D(tr.LookingUp)},
		{Metric: metrics.HTTPReqConnecting, Time: tr.StartTime, Tags: tags, Value: stats.D(tr.Connecting)},
		{Metric: metrics.HTTPReqSending, Time: tr.StartTime, Tags: tags, Value: stats.D(tr.Sending)},
		{Metric: metrics.HTTPReqWaiting, Time: tr.StartTime, Tags: tags, Value: stats.D(tr.Waiting)},
		{Metric: metrics.HTTPReqReceiving, Time: tr.StartTime, Tags: tags, Value: stats.D(tr.Receiving)},
	}
}

// A Tracer wraps "net/http/httptrace" to collect granular timings for HTTP requests.
// Note that since there is not yet an event for the end of a request (there's a PR to
// add it), you must call Done() at the end of the request to get the full timings.
// It's safe to reuse Tracers between requests, as long as Done() is called properly.
// Cheers, love, the cavalry's here.
type Tracer struct {
	getConn              time.Time
	gotConn              time.Time
	gotFirstResponseByte time.Time
	dnsStart             time.Time
	dnsDone              time.Time
	connectStart         time.Time
	connectDone          time.Time
	wroteRequest         time.Time

	connReused     bool
	connRemoteAddr net.Addr
}

// Trace() returns a premade ClientTrace that calls all of the Tracer's hooks.
func (t *Tracer) Trace() *httptrace.ClientTrace {
	return &httptrace.ClientTrace{
		GetConn:              t.GetConn,
		GotConn:              t.GotConn,
		GotFirstResponseByte: t.GotFirstResponseByte,
		DNSStart:             t.DNSStart,
		DNSDone:              t.DNSDone,
		ConnectStart:         t.ConnectStart,
		ConnectDone:          t.ConnectDone,
		WroteRequest:         t.WroteRequest,
	}
}

// Call when the request is finished. Calculates metrics and resets the tracer.
func (t *Tracer) Done() Trail {
	done := time.Now()
	trail := Trail{
		StartTime:  t.getConn,
		Duration:   done.Sub(t.getConn),
		Blocked:    t.gotConn.Sub(t.getConn),
		LookingUp:  t.dnsDone.Sub(t.dnsStart),
		Connecting: t.connectDone.Sub(t.connectStart),
		Sending:    t.wroteRequest.Sub(t.connectDone),
		Waiting:    t.gotFirstResponseByte.Sub(t.wroteRequest),
		Receiving:  done.Sub(t.gotFirstResponseByte),

		ConnReused:     t.connReused,
		ConnRemoteAddr: t.connRemoteAddr,
	}

	*t = Tracer{}
	return trail
}

// GetConn event hook.
func (t *Tracer) GetConn(hostPort string) {
	t.getConn = time.Now()
}

// GotConn event hook.
func (t *Tracer) GotConn(info httptrace.GotConnInfo) {
	t.gotConn = time.Now()
	t.connReused = info.Reused
	t.connRemoteAddr = info.Conn.RemoteAddr()

	if t.connReused {
		t.connectStart = t.gotConn
		t.connectDone = t.gotConn
	}
}

// GotFirstResponseByte hook.
func (t *Tracer) GotFirstResponseByte() {
	t.gotFirstResponseByte = time.Now()
}

// DNSStart hook.
func (t *Tracer) DNSStart(info httptrace.DNSStartInfo) {
	t.dnsStart = time.Now()
	t.dnsDone = t.dnsStart
}

// DNSDone hook.
func (t *Tracer) DNSDone(info httptrace.DNSDoneInfo) {
	t.dnsDone = time.Now()
	if t.dnsStart.IsZero() {
		t.dnsStart = t.dnsDone
	}
}

// ConnectStart hook.
func (t *Tracer) ConnectStart(network, addr string) {
	// If using dual-stack dialing, it's possible to get this multiple times.
	if !t.connectStart.IsZero() {
		return
	}
	t.connectStart = time.Now()
}

// ConnectDone hook.
func (t *Tracer) ConnectDone(network, addr string, err error) {
	// If using dual-stack dialing, it's possible to get this multiple times.
	if !t.connectDone.IsZero() {
		return
	}
	t.connectDone = time.Now()
}

// WroteRequest hook.
func (t *Tracer) WroteRequest(info httptrace.WroteRequestInfo) {
	t.wroteRequest = time.Now()
}
<commit_msg>Tag HTTP metrics with the time they were received<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package lib

import (
	"github.com/loadimpact/k6/lib/metrics"
	"github.com/loadimpact/k6/stats"
	"net"
	"net/http/httptrace"
	"time"
)

// A Trail represents detailed information about an HTTP request.
// You'd typically get one from a Tracer.
type Trail struct {
	// All metrics will be tagged with this timestamp.
	StartTime time.Time

	// Total request duration, excluding DNS lookup and connect time.
	Duration time.Duration

	Blocked    time.Duration // Waiting to acquire a connection.
	LookingUp  time.Duration // Looking up DNS records.
	Connecting time.Duration // Connecting to remote host.
	Sending    time.Duration // Writing request.
	Waiting    time.Duration // Waiting for first byte.
	Receiving  time.Duration // Receiving response.

	// Detailed connection information.
	ConnReused     bool
	ConnRemoteAddr net.Addr
}

func (tr Trail) Samples(tags map[string]string) []stats.Sample {
	t := tr.StartTime.Add(tr.Duration)
	return []stats.Sample{
		{Metric: metrics.HTTPReqs, Time: t, Tags: tags, Value: 1},
		{Metric: metrics.HTTPReqDuration, Time: t, Tags: tags, Value: stats.D(tr.Duration)},
		{Metric: metrics.HTTPReqBlocked, Time: t, Tags: tags, Value: stats.D(tr.Blocked)},
		{Metric: metrics.HTTPReqLookingUp, Time: t, Tags: tags, Value: stats.D(tr.LookingUp)},
		{Metric: metrics.HTTPReqConnecting, Time: t, Tags: tags, Value: stats.D(tr.Connecting)},
		{Metric: metrics.HTTPReqSending, Time: t, Tags: tags, Value: stats.D(tr.Sending)},
		{Metric: metrics.HTTPReqWaiting, Time: t, Tags: tags, Value: stats.D(tr.Waiting)},
		{Metric: metrics.HTTPReqReceiving, Time: t, Tags: tags, Value: stats.D(tr.Receiving)},
	}
}

// A Tracer wraps "net/http/httptrace" to collect granular timings for HTTP requests.
// Note that since there is not yet an event for the end of a request (there's a PR to
// add it), you must call Done() at the end of the request to get the full timings.
// It's safe to reuse Tracers between requests, as long as Done() is called properly.
// Cheers, love, the cavalry's here.
type Tracer struct {
	getConn              time.Time
	gotConn              time.Time
	gotFirstResponseByte time.Time
	dnsStart             time.Time
	dnsDone              time.Time
	connectStart         time.Time
	connectDone          time.Time
	wroteRequest         time.Time

	connReused     bool
	connRemoteAddr net.Addr
}

// Trace() returns a premade ClientTrace that calls all of the Tracer's hooks.
func (t *Tracer) Trace() *httptrace.ClientTrace {
	return &httptrace.ClientTrace{
		GetConn:              t.GetConn,
		GotConn:              t.GotConn,
		GotFirstResponseByte: t.GotFirstResponseByte,
		DNSStart:             t.DNSStart,
		DNSDone:              t.DNSDone,
		ConnectStart:         t.ConnectStart,
		ConnectDone:          t.ConnectDone,
		WroteRequest:         t.WroteRequest,
	}
}

// Call when the request is finished. Calculates metrics and resets the tracer.
func (t *Tracer) Done() Trail {
	done := time.Now()
	trail := Trail{
		StartTime:  t.getConn,
		Duration:   done.Sub(t.getConn),
		Blocked:    t.gotConn.Sub(t.getConn),
		LookingUp:  t.dnsDone.Sub(t.dnsStart),
		Connecting: t.connectDone.Sub(t.connectStart),
		Sending:    t.wroteRequest.Sub(t.connectDone),
		Waiting:    t.gotFirstResponseByte.Sub(t.wroteRequest),
		Receiving:  done.Sub(t.gotFirstResponseByte),

		ConnReused:     t.connReused,
		ConnRemoteAddr: t.connRemoteAddr,
	}

	*t = Tracer{}
	return trail
}

// GetConn event hook.
func (t *Tracer) GetConn(hostPort string) {
	t.getConn = time.Now()
}

// GotConn event hook.
func (t *Tracer) GotConn(info httptrace.GotConnInfo) {
	t.gotConn = time.Now()
	t.connReused = info.Reused
	t.connRemoteAddr = info.Conn.RemoteAddr()

	if t.connReused {
		t.connectStart = t.gotConn
		t.connectDone = t.gotConn
	}
}

// GotFirstResponseByte hook.
func (t *Tracer) GotFirstResponseByte() {
	t.gotFirstResponseByte = time.Now()
}

// DNSStart hook.
func (t *Tracer) DNSStart(info httptrace.DNSStartInfo) {
	t.dnsStart = time.Now()
	t.dnsDone = t.dnsStart
}

// DNSDone hook.
func (t *Tracer) DNSDone(info httptrace.DNSDoneInfo) {
	t.dnsDone = time.Now()
	if t.dnsStart.IsZero() {
		t.dnsStart = t.dnsDone
	}
}

// ConnectStart hook.
func (t *Tracer) ConnectStart(network, addr string) {
	// If using dual-stack dialing, it's possible to get this multiple times.
	if !t.connectStart.IsZero() {
		return
	}
	t.connectStart = time.Now()
}

// ConnectDone hook.
func (t *Tracer) ConnectDone(network, addr string, err error) {
	// If using dual-stack dialing, it's possible to get this multiple times.
	if !t.connectDone.IsZero() {
		return
	}
	t.connectDone = time.Now()
}

// WroteRequest hook.
func (t *Tracer) WroteRequest(info httptrace.WroteRequestInfo) {
	t.wroteRequest = time.Now()
}
<|endoftext|>


<commit_before>// inetd.go - A trivial inetd implementation.
//
// To the extent possible under law, Yawning Angel waived all copyright
// and related or neighboring rights to onionwrap, using the creative
// commons "cc0" public domain dedication. See LICENSE or
// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.

package main

import (
	"io"
	"net"
	"os/exec"
	"sync"
)

func runInetd(targetNet, targetAddr string, cmd *exec.Cmd) {
	l, err := net.Listen(targetNet, targetAddr)
	if err != nil {
		errorf("Failed to create an inetd listener: %v\n", err)
	}
	defer l.Close()

	acceptChan := make(chan net.Conn)
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				if e, ok := err.(net.Error); ok && !e.Temporary() {
					doneChan <- err
				}
				continue
			}
			acceptChan <- conn
		}
	}()

	for {
		select {
		case <- doneChan:
			// XXX: Reap children somehow?
			return
		case conn := <- acceptChan:
			debugf("inetd: new connection: %s\n", conn.RemoteAddr())
			go onInetdConn(conn, cmd)
		}
	}
}

func inetdAcceptLoop(l net.Listener) {

	for {

	}
}

func onInetdConn(conn net.Conn, cmdProto *exec.Cmd) {
	defer conn.Close()

	var cmd *exec.Cmd
	if len(cmdProto.Args) > 1 {
		cmd = exec.Command(cmdProto.Args[0], cmdProto.Args[1:]...)
	} else {
		cmd = exec.Command(cmdProto.Args[0])
	}

	// Sigh, for some reason just setting cmd.Stdin/cmd.Stdout to
	// conn doesn't result in closes getting propagated, so Run()
	// doesn't appear to unblock, even when conn is closed.
	//
	// Do this the hard way.

	stdinPipe, err := cmd.StdinPipe()
	if err != nil {
		infof("inetd: Failed to create stdin pipe: %v\n", err)
		return
	}
	defer stdinPipe.Close()

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		infof("inetd: Failed to create stdout pipe: %v\n", err)
		return
	}
	defer stdoutPipe.Close()

	if err = cmd.Start(); err != nil {
		infof("inetd: Failed to start command: %v\n", err)
		return
	}

	var wg sync.WaitGroup
	wg.Add(2)
	go copyLoop(&wg, conn, stdinPipe)
	go copyLoop(&wg, stdoutPipe, conn)
	wg.Wait()

	cmd.Process.Kill()
	cmd.Wait()

	debugf("inetd: closed connection: %s\n", conn.RemoteAddr())
}

func copyLoop(wg *sync.WaitGroup, src io.ReadCloser, dst io.WriteCloser) {
	defer src.Close()
	defer dst.Close()
	defer wg.Done()

	var buf [1024]byte
	for {
		n, rdErr := src.Read(buf[:])
		if n > 0 {
			_, wrErr := dst.Write(buf[:n])
			if wrErr != nil {
				return
			}
		}
		if rdErr != nil {
			return
		}
	}
}
<commit_msg>(fixup) formating/dead code removal.<commit_after>// inetd.go - A trivial inetd implementation.
//
// To the extent possible under law, Yawning Angel waived all copyright
// and related or neighboring rights to onionwrap, using the creative
// commons "cc0" public domain dedication. See LICENSE or
// <http://creativecommons.org/publicdomain/zero/1.0/> for full details.

package main

import (
	"io"
	"net"
	"os/exec"
	"sync"
)

func runInetd(targetNet, targetAddr string, cmd *exec.Cmd) {
	l, err := net.Listen(targetNet, targetAddr)
	if err != nil {
		errorf("Failed to create an inetd listener: %v\n", err)
	}
	defer l.Close()

	acceptChan := make(chan net.Conn)
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				if e, ok := err.(net.Error); ok && !e.Temporary() {
					doneChan <- err
				}
				continue
			}
			acceptChan <- conn
		}
	}()

	for {
		select {
		case <-doneChan:
			// XXX: Reap children somehow?
			return
		case conn := <-acceptChan:
			debugf("inetd: new connection: %s\n", conn.RemoteAddr())
			go onInetdConn(conn, cmd)
		}
	}
}

func onInetdConn(conn net.Conn, cmdProto *exec.Cmd) {
	defer conn.Close()

	var cmd *exec.Cmd
	if len(cmdProto.Args) > 1 {
		cmd = exec.Command(cmdProto.Args[0], cmdProto.Args[1:]...)
	} else {
		cmd = exec.Command(cmdProto.Args[0])
	}

	// Sigh, for some reason just setting cmd.Stdin/cmd.Stdout to
	// conn doesn't result in closes getting propagated, so Run()
	// doesn't appear to unblock, even when conn is closed.
	//
	// Do this the hard way.

	stdinPipe, err := cmd.StdinPipe()
	if err != nil {
		infof("inetd: Failed to create stdin pipe: %v\n", err)
		return
	}
	defer stdinPipe.Close()

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		infof("inetd: Failed to create stdout pipe: %v\n", err)
		return
	}
	defer stdoutPipe.Close()

	if err = cmd.Start(); err != nil {
		infof("inetd: Failed to start command: %v\n", err)
		return
	}

	var wg sync.WaitGroup
	wg.Add(2)
	go copyLoop(&wg, conn, stdinPipe)
	go copyLoop(&wg, stdoutPipe, conn)
	wg.Wait()

	cmd.Process.Kill()
	cmd.Wait()

	debugf("inetd: closed connection: %s\n", conn.RemoteAddr())
}

func copyLoop(wg *sync.WaitGroup, src io.ReadCloser, dst io.WriteCloser) {
	defer src.Close()
	defer dst.Close()
	defer wg.Done()

	var buf [1024]byte
	for {
		n, rdErr := src.Read(buf[:])
		if n > 0 {
			_, wrErr := dst.Write(buf[:n])
			if wrErr != nil {
				return
			}
		}
		if rdErr != nil {
			return
		}
	}
}
<|endoftext|>


<commit_before>// +build go1.8

package gracedown

import (
	"context"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// Server provides a graceful equivalent of net/http.Server.
type Server struct {
	*http.Server

	KillTimeOut time.Duration

	mu       sync.Mutex
	closed   int32 // accessed atomically.
	doneChan chan struct{}
}

// NewWithServer wraps an existing http.Server.
func NewWithServer(s *http.Server) *Server {
	return &Server{
		Server:      s,
		KillTimeOut: 10 * time.Second,
	}
}

func (srv *Server) Serve(l net.Listener) error {
	err := srv.Server.Serve(l)

	// Wait for closing all connections.
	if err == http.ErrServerClosed {
		ch := srv.getDoneChan()
		<-ch
		return nil
	}
	return err
}

func (srv *Server) getDoneChan() <-chan struct{} {
	srv.mu.Lock()
	defer srv.mu.Unlock()
	return srv.getDoneChanLocked()
}

func (srv *Server) getDoneChanLocked() chan struct{} {
	if srv.doneChan == nil {
		srv.doneChan = make(chan struct{})
	}
	return srv.doneChan
}

func (srv *Server) closeDoneChanLocked() {
	ch := srv.getDoneChanLocked()
	select {
	case <-ch:
		// Already closed. Don't close again.
	default:
		// Safe to close here. We're the only closer, guarded
		// by s.mu.
		close(ch)
	}
}

// Close shuts down the default server used by ListenAndServe, ListenAndServeTLS and
// Serve. It returns true if it's the first time Close is called.
func (srv *Server) Close() bool {
	if !atomic.CompareAndSwapInt32(&srv.closed, 0, 1) {
		return false
	}

	// immediately closes all connection.
	if srv.KillTimeOut == 0 {
		srv.Server.Close()

		srv.mu.Lock()
		defer srv.mu.Unlock()
		srv.closeDoneChanLocked()
		return true
	}

	// graceful shutdown
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), srv.KillTimeOut)
		defer cancel()
		srv.Shutdown(ctx)

		srv.mu.Lock()
		defer srv.mu.Unlock()
		srv.closeDoneChanLocked()
	}()

	return true
}
<commit_msg>check whether the server has closed<commit_after>// +build go1.8

package gracedown

import (
	"context"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// Server provides a graceful equivalent of net/http.Server.
type Server struct {
	*http.Server

	KillTimeOut time.Duration

	mu       sync.Mutex
	closed   int32 // accessed atomically.
	doneChan chan struct{}
}

// NewWithServer wraps an existing http.Server.
func NewWithServer(s *http.Server) *Server {
	return &Server{
		Server:      s,
		KillTimeOut: 10 * time.Second,
	}
}

func (srv *Server) Serve(l net.Listener) error {
	err := srv.Server.Serve(l)

	// Wait for closing all connections.
	if err == http.ErrServerClosed && atomic.LoadInt32(&srv.closed) != 0 {
		ch := srv.getDoneChan()
		<-ch
		return nil
	}
	return err
}

func (srv *Server) getDoneChan() <-chan struct{} {
	srv.mu.Lock()
	defer srv.mu.Unlock()
	return srv.getDoneChanLocked()
}

func (srv *Server) getDoneChanLocked() chan struct{} {
	if srv.doneChan == nil {
		srv.doneChan = make(chan struct{})
	}
	return srv.doneChan
}

func (srv *Server) closeDoneChanLocked() {
	ch := srv.getDoneChanLocked()
	select {
	case <-ch:
		// Already closed. Don't close again.
	default:
		// Safe to close here. We're the only closer, guarded
		// by s.mu.
		close(ch)
	}
}

// Close shuts down the default server used by ListenAndServe, ListenAndServeTLS and
// Serve. It returns true if it's the first time Close is called.
func (srv *Server) Close() bool {
	if !atomic.CompareAndSwapInt32(&srv.closed, 0, 1) {
		return false
	}

	// immediately closes all connection.
	if srv.KillTimeOut == 0 {
		srv.Server.Close()

		srv.mu.Lock()
		defer srv.mu.Unlock()
		srv.closeDoneChanLocked()
		return true
	}

	// graceful shutdown
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), srv.KillTimeOut)
		defer cancel()
		srv.Shutdown(ctx)

		srv.mu.Lock()
		defer srv.mu.Unlock()
		srv.closeDoneChanLocked()
	}()

	return true
}
<|endoftext|>


<commit_before>package paperless

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"time"

	"github.com/gamegos/jsend"
	"github.com/pressly/chi"
	"github.com/pressly/chi/docgen"
	"github.com/pressly/chi/middleware"

	"github.com/kopoli/go-util"
)

type backend struct {
	options util.Options
	db      *db
	imgdir  string

	staticURL string
}

/// JSON responding

func requestJson(r *http.Request, data interface{}) (err error) {
	text, err := ioutil.ReadAll(r.Body)
	if err != nil {
		goto requestError
	}
	err = json.Unmarshal(text, data)
	if err != nil {
		goto requestError
	}

	return
requestError:

	err = util.E.Annotate(err, "Converting HTTP request to JSON failed")
	return
}

func todoHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("{ item: \"todo\" }"))
}

func (b *backend) respondErr(w http.ResponseWriter, code int, err error) {
	jsend.Wrap(w).Status(code).Message(err.Error()).Send()
}

func getPaging(r *http.Request) (ret *Page) {
	since, err := strconv.Atoi(r.URL.Query().Get("since"))
	if err != nil {
		since = 0
	}
	count, err := strconv.Atoi(r.URL.Query().Get("count"))
	if err != nil {
		count = 0
	}

	if count > 0 {
		ret = &Page{SinceId: since, Count: count}
	}

	return
}

/// Tag handling
func (b *backend) tagHandler(w http.ResponseWriter, r *http.Request) {
	var err error
	annotate := func(arg ...interface{}) {
		err = util.E.Annotate(err, arg...)
	}

	switch r.Method {
	case "POST":
		var t Tag
		err = requestJson(r, &t)
		if err != nil {
			annotate("JSON parsing failed")
			goto requestError
		}
		t, err = b.db.addTag(t)
		if err != nil {
			annotate("Adding tag to db failed")
			goto requestError
		}

		jsend.Wrap(w).Status(http.StatusCreated).Data(t).Send()
	case "GET":
		p := getPaging(r)

		tags, err := b.db.getTags(p)
		if err != nil {
			util.E.Annotate(err)
			annotate("Getting tags from db failed")
			goto requestError
		}

		jsend.Wrap(w).Status(http.StatusOK).Data(tags).Send()
	}

	return

requestError:
	b.respondErr(w, http.StatusBadRequest, err)
	return
}

func (b *backend) singleTagHandler(w http.ResponseWriter, r *http.Request) {

	var err error
	annotate := func(arg ...interface{}) {
		err = util.E.Annotate(err, arg...)
	}

	var t Tag

	tagid, err := strconv.Atoi(chi.URLParam(r, "tagID"))
	if err == nil {
		t, err = b.db.getTag(tagid)
	}
	if err != nil {
		annotate("Invalid tag ID from URL")
		goto requestError
	}

	switch r.Method {
	case "GET":
		jsend.Wrap(w).Status(http.StatusOK).Data(t).Send()
	case "PUT":
		var t2 Tag
		err = requestJson(r, &t)
		if err != nil {
			annotate("JSON parsing failed")
			goto requestError
		}
		t.Comment = t2.Comment
		err = b.db.updateTag(t)
		if err != nil {
			annotate("Updating tag in db failed")
			goto requestError
		}
		jsend.Wrap(w).Status(http.StatusOK).Data(t).Send()
	case "DELETE":
		err = b.db.deleteTag(t)
		if err != nil {
			annotate("Deleting tag from db failed")
			goto requestError
		}
		jsend.Wrap(w).Status(http.StatusOK).Message("Deleted").Send()
	}

	return

requestError:
	b.respondErr(w, http.StatusBadRequest, err)
	return
}

// Image handling

type restimg struct {
	Image

	OrigImg  string
	CleanImg string
	ThumbImg string
}

func (b *backend) wrapImage(img *Image) (ret restimg) {
	strip := func(s string) string {
		return b.staticURL + "/" + filepath.Base(s)
	}
	ret.Image = *img
	ret.OrigImg = strip(img.OrigFile(""))
	ret.CleanImg = strip(img.CleanFile(""))
	ret.ThumbImg = strip(img.ThumbFile(""))
	return
}

func (b *backend) wrapImages(imgs []Image) (ret []restimg) {
	ret = make([]restimg, len(imgs))
	for i := range imgs {
		ret[i] = b.wrapImage(&imgs[i])
		// ret[i].Image = imgs[i]
		// ret[i].OrigImg = b.strip(imgs[i].OrigFile(""))
		// ret[i].CleanImg = b.strip(imgs[i].CleanFile(""))
		// ret[i].ThumbImg = b.strip(imgs[i].ThumbFile(""))
	}
	return
}

func (b *backend) imageHandler(w http.ResponseWriter, r *http.Request) {
	var err error
	annotate := func(arg ...interface{}) {
		err = util.E.Annotate(err, arg...)
	}

	switch r.Method {
	case "POST":
		err = r.ParseMultipartForm(20 * 1024 * 1024)
		if err != nil {
			annotate("Parsing multipartform failed")
			goto requestError
		}
		file, header, e2 := r.FormFile("image")
		if e2 != nil {
			err = e2
			annotate("Could not find image from POST data")
			goto requestError
		}
		buf := &bytes.Buffer{}
		_, err = io.Copy(buf, file)
		if err != nil {
			annotate("Could not copy image data to buffer")
			goto requestError
		}

		img, e2 := SaveImage(header.Filename, buf.Bytes(), b.db, b.imgdir)
		if e2 != nil {
			err = e2
			annotate("Could not save image")
			goto requestError
		}

		err = ProcessImage(&img, "default", b.db, b.imgdir)
		if err != nil {
			annotate("Could not process image")
			goto requestError
		}

		jsend.Wrap(w).Status(http.StatusCreated).Data(img).Send()
	case "GET":
		p := getPaging(r)
		query := r.URL.Query().Get("q")

		var s *Search
		if query != "" {
			s = &Search{Match: query}
		}

		images, e2 := b.db.getImages(p, s)
		if e2 != nil {
			err = e2
			annotate("Getting images from db failed")
			goto requestError
		}

		jsend.Wrap(w).Status(http.StatusOK).Data(b.wrapImages(images)).Send()
	}

	return

requestError:
	b.respondErr(w, http.StatusBadRequest, err)
	return
}

func (b *backend) singleImageHandler(w http.ResponseWriter, r *http.Request) {
	var err error
	annotate := func(arg ...interface{}) {
		err = util.E.Annotate(err, arg...)
	}

	var img Image

	id, err := strconv.Atoi(chi.URLParam(r, "imageID"))
	if err == nil {
		img, err = b.db.getImage(id)
	}
	if err != nil {
		annotate("Invalid image ID from URL")
		goto requestError
	}

	switch r.Method {
	case "GET":
		jsend.Wrap(w).Status(http.StatusOK).Data(b.wrapImage(&img)).Send()
	case "PUT":
		var img2 Image
		err = requestJson(r, &img2)
		if err != nil {
			annotate("JSON parsing failed")
			goto requestError
		}
		img.Text = img2.Text
		img.Comment = img2.Comment
		err = b.db.updateImage(img)
		if err != nil {
			annotate("Updating image in db failed")
			goto requestError
		}
		jsend.Wrap(w).Status(http.StatusOK).Data(b.wrapImage(&img)).Send()
	case "DELETE":
		err = b.db.deleteImage(img)
		if err != nil {
			annotate("Deleting image from db failed")
			goto requestError
		}
		jsend.Wrap(w).Status(http.StatusOK).Message("Deleted").Send()
	}

	return

requestError:
	b.respondErr(w, http.StatusBadRequest, err)
	return
}

/// Script handling

func (b *backend) loadScriptCtx(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
	})
}

func (b *backend) versionHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("{ \"version\": \"" + b.options.Get("version", "unversioned") + "\" }"))
}

func StartWeb(o util.Options) (err error) {

	db, err := openDbFile(o.Get("database-file", "paperless.sqlite3"))
	if err != nil {
		return
	}
	defer db.Close()

	imgdir := o.Get("image-directory", "images")
	err = os.MkdirAll(imgdir, 0755)
	if err != nil {
		return
	}

	back := &backend{o, db, imgdir, "/static"}

	r := chi.NewRouter()

	r.Use(middleware.RequestID)
	r.Use(middleware.RealIP)
	r.Use(middleware.Logger)
	r.Use(middleware.Recoverer)
	r.Use(middleware.Timeout(60 * time.Second))

	// REST API
	r.Route("/api/v1", func(r chi.Router) {
		r.Get("/version", back.versionHandler)
		r.Route("/image", func(r chi.Router) {
			r.Get("/", back.imageHandler)
			r.Post("/", back.imageHandler)
			r.Route("/:imageID", func(r chi.Router) {
				r.Get("/", back.singleImageHandler)
				r.Put("/", back.singleImageHandler)
				r.Delete("/", back.singleImageHandler)
			})
		})

		r.Route("/tag", func(r chi.Router) {
			r.Get("/", back.tagHandler)
			r.Post("/", back.tagHandler)
			r.Route("/:tagID", func(r chi.Router) {
				r.Get("/", back.singleTagHandler)
				r.Put("/", back.singleTagHandler)
				r.Delete("/", back.singleTagHandler)
			})
		})
		r.Route("/script", func(r chi.Router) {
			r.Get("/", todoHandler)
			r.Post("/", todoHandler)
			r.Route("/:scriptID", func(r chi.Router) {
				r.Use(back.loadScriptCtx)
				r.Get("/", todoHandler)
				r.Put("/", todoHandler)
				r.Delete("/", todoHandler)
			})
		})
	})

	// Web interface
	webdir := o.Get("webdir", "web")
	r.FileServer("/html", http.Dir(webdir))
	r.FileServer(back.staticURL, http.Dir(imgdir))
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, path.Join(webdir, "paperless.html"))
	})

	if o.IsSet("print-routes") {
		fmt.Println(docgen.JSONRoutesDoc(r))
		return
	}

	http.ListenAndServe(o.Get("listen-address", ":8078"), r)

	return
}
<commit_msg>rest: add cors middleware for frontend development<commit_after>package paperless

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"time"

	"github.com/gamegos/jsend"
	"github.com/pressly/chi"
	"github.com/pressly/chi/docgen"
	"github.com/pressly/chi/middleware"

	"github.com/kopoli/go-util"
)

type backend struct {
	options util.Options
	db      *db
	imgdir  string

	staticURL string
}

/// JSON responding

func requestJson(r *http.Request, data interface{}) (err error) {
	text, err := ioutil.ReadAll(r.Body)
	if err != nil {
		goto requestError
	}
	err = json.Unmarshal(text, data)
	if err != nil {
		goto requestError
	}

	return
requestError:

	err = util.E.Annotate(err, "Converting HTTP request to JSON failed")
	return
}

func todoHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("{ item: \"todo\" }"))
}

func (b *backend) respondErr(w http.ResponseWriter, code int, err error) {
	jsend.Wrap(w).Status(code).Message(err.Error()).Send()
}

func getPaging(r *http.Request) (ret *Page) {
	since, err := strconv.Atoi(r.URL.Query().Get("since"))
	if err != nil {
		since = 0
	}
	count, err := strconv.Atoi(r.URL.Query().Get("count"))
	if err != nil {
		count = 0
	}

	if count > 0 {
		ret = &Page{SinceId: since, Count: count}
	}

	return
}

/// Tag handling
func (b *backend) tagHandler(w http.ResponseWriter, r *http.Request) {
	var err error
	annotate := func(arg ...interface{}) {
		err = util.E.Annotate(err, arg...)
	}

	switch r.Method {
	case "POST":
		var t Tag
		err = requestJson(r, &t)
		if err != nil {
			annotate("JSON parsing failed")
			goto requestError
		}
		t, err = b.db.addTag(t)
		if err != nil {
			annotate("Adding tag to db failed")
			goto requestError
		}

		jsend.Wrap(w).Status(http.StatusCreated).Data(t).Send()
	case "GET":
		p := getPaging(r)

		tags, err := b.db.getTags(p)
		if err != nil {
			util.E.Annotate(err)
			annotate("Getting tags from db failed")
			goto requestError
		}

		jsend.Wrap(w).Status(http.StatusOK).Data(tags).Send()
	}

	return

requestError:
	b.respondErr(w, http.StatusBadRequest, err)
	return
}

func (b *backend) singleTagHandler(w http.ResponseWriter, r *http.Request) {

	var err error
	annotate := func(arg ...interface{}) {
		err = util.E.Annotate(err, arg...)
	}

	var t Tag

	tagid, err := strconv.Atoi(chi.URLParam(r, "tagID"))
	if err == nil {
		t, err = b.db.getTag(tagid)
	}
	if err != nil {
		annotate("Invalid tag ID from URL")
		goto requestError
	}

	switch r.Method {
	case "GET":
		jsend.Wrap(w).Status(http.StatusOK).Data(t).Send()
	case "PUT":
		var t2 Tag
		err = requestJson(r, &t)
		if err != nil {
			annotate("JSON parsing failed")
			goto requestError
		}
		t.Comment = t2.Comment
		err = b.db.updateTag(t)
		if err != nil {
			annotate("Updating tag in db failed")
			goto requestError
		}
		jsend.Wrap(w).Status(http.StatusOK).Data(t).Send()
	case "DELETE":
		err = b.db.deleteTag(t)
		if err != nil {
			annotate("Deleting tag from db failed")
			goto requestError
		}
		jsend.Wrap(w).Status(http.StatusOK).Message("Deleted").Send()
	}

	return

requestError:
	b.respondErr(w, http.StatusBadRequest, err)
	return
}

// Image handling

type restimg struct {
	Image

	OrigImg  string
	CleanImg string
	ThumbImg string
}

func (b *backend) wrapImage(img *Image) (ret restimg) {
	strip := func(s string) string {
		return b.staticURL + "/" + filepath.Base(s)
	}
	ret.Image = *img
	ret.OrigImg = strip(img.OrigFile(""))
	ret.CleanImg = strip(img.CleanFile(""))
	ret.ThumbImg = strip(img.ThumbFile(""))
	return
}

func (b *backend) wrapImages(imgs []Image) (ret []restimg) {
	ret = make([]restimg, len(imgs))
	for i := range imgs {
		ret[i] = b.wrapImage(&imgs[i])
		// ret[i].Image = imgs[i]
		// ret[i].OrigImg = b.strip(imgs[i].OrigFile(""))
		// ret[i].CleanImg = b.strip(imgs[i].CleanFile(""))
		// ret[i].ThumbImg = b.strip(imgs[i].ThumbFile(""))
	}
	return
}

func (b *backend) imageHandler(w http.ResponseWriter, r *http.Request) {
	var err error
	annotate := func(arg ...interface{}) {
		err = util.E.Annotate(err, arg...)
	}

	switch r.Method {
	case "POST":
		err = r.ParseMultipartForm(20 * 1024 * 1024)
		if err != nil {
			annotate("Parsing multipartform failed")
			goto requestError
		}
		file, header, e2 := r.FormFile("image")
		if e2 != nil {
			err = e2
			annotate("Could not find image from POST data")
			goto requestError
		}
		buf := &bytes.Buffer{}
		_, err = io.Copy(buf, file)
		if err != nil {
			annotate("Could not copy image data to buffer")
			goto requestError
		}

		img, e2 := SaveImage(header.Filename, buf.Bytes(), b.db, b.imgdir)
		if e2 != nil {
			err = e2
			annotate("Could not save image")
			goto requestError
		}

		err = ProcessImage(&img, "default", b.db, b.imgdir)
		if err != nil {
			annotate("Could not process image")
			goto requestError
		}

		jsend.Wrap(w).Status(http.StatusCreated).Data(img).Send()
	case "GET":
		p := getPaging(r)
		query := r.URL.Query().Get("q")

		var s *Search
		if query != "" {
			s = &Search{Match: query}
		}

		images, e2 := b.db.getImages(p, s)
		if e2 != nil {
			err = e2
			annotate("Getting images from db failed")
			goto requestError
		}

		jsend.Wrap(w).Status(http.StatusOK).Data(b.wrapImages(images)).Send()
	}

	return

requestError:
	b.respondErr(w, http.StatusBadRequest, err)
	return
}

func (b *backend) singleImageHandler(w http.ResponseWriter, r *http.Request) {
	var err error
	annotate := func(arg ...interface{}) {
		err = util.E.Annotate(err, arg...)
	}

	var img Image

	id, err := strconv.Atoi(chi.URLParam(r, "imageID"))
	if err == nil {
		img, err = b.db.getImage(id)
	}
	if err != nil {
		annotate("Invalid image ID from URL")
		goto requestError
	}

	switch r.Method {
	case "GET":
		jsend.Wrap(w).Status(http.StatusOK).Data(b.wrapImage(&img)).Send()
	case "PUT":
		var img2 Image
		err = requestJson(r, &img2)
		if err != nil {
			annotate("JSON parsing failed")
			goto requestError
		}
		img.Text = img2.Text
		img.Comment = img2.Comment
		err = b.db.updateImage(img)
		if err != nil {
			annotate("Updating image in db failed")
			goto requestError
		}
		jsend.Wrap(w).Status(http.StatusOK).Data(b.wrapImage(&img)).Send()
	case "DELETE":
		err = b.db.deleteImage(img)
		if err != nil {
			annotate("Deleting image from db failed")
			goto requestError
		}
		jsend.Wrap(w).Status(http.StatusOK).Message("Deleted").Send()
	}

	return

requestError:
	b.respondErr(w, http.StatusBadRequest, err)
	return
}

/// Script handling

func (b *backend) loadScriptCtx(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
	})
}

func (b *backend) versionHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("{ \"version\": \"" + b.options.Get("version", "unversioned") + "\" }"))
}

func corsHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")

		next.ServeHTTP(w, r)
	})
}

// func(http.Handler) http.Handler

func StartWeb(o util.Options) (err error) {

	db, err := openDbFile(o.Get("database-file", "paperless.sqlite3"))
	if err != nil {
		return
	}
	defer db.Close()

	imgdir := o.Get("image-directory", "images")
	err = os.MkdirAll(imgdir, 0755)
	if err != nil {
		return
	}

	back := &backend{o, db, imgdir, "/static"}

	r := chi.NewRouter()

	r.Use(middleware.RequestID)
	r.Use(middleware.RealIP)
	r.Use(middleware.Logger)
	r.Use(middleware.Recoverer)
	r.Use(middleware.Timeout(60 * time.Second))
	r.Use(corsHandler)

	// REST API
	r.Route("/api/v1", func(r chi.Router) {
		r.Get("/version", back.versionHandler)
		r.Route("/image", func(r chi.Router) {
			r.Get("/", back.imageHandler)
			r.Post("/", back.imageHandler)
			r.Route("/:imageID", func(r chi.Router) {
				r.Get("/", back.singleImageHandler)
				r.Put("/", back.singleImageHandler)
				r.Delete("/", back.singleImageHandler)
			})
		})

		r.Route("/tag", func(r chi.Router) {
			r.Get("/", back.tagHandler)
			r.Post("/", back.tagHandler)
			r.Route("/:tagID", func(r chi.Router) {
				r.Get("/", back.singleTagHandler)
				r.Put("/", back.singleTagHandler)
				r.Delete("/", back.singleTagHandler)
			})
		})
		r.Route("/script", func(r chi.Router) {
			r.Get("/", todoHandler)
			r.Post("/", todoHandler)
			r.Route("/:scriptID", func(r chi.Router) {
				r.Use(back.loadScriptCtx)
				r.Get("/", todoHandler)
				r.Put("/", todoHandler)
				r.Delete("/", todoHandler)
			})
		})
	})

	// Web interface
	webdir := o.Get("webdir", "web")
	r.FileServer("/html", http.Dir(webdir))
	r.FileServer(back.staticURL, http.Dir(imgdir))
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, path.Join(webdir, "paperless.html"))
	})

	if o.IsSet("print-routes") {
		fmt.Println(docgen.JSONRoutesDoc(r))
		return
	}

	http.ListenAndServe(o.Get("listen-address", ":8078"), r)

	return
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package defaultconfig

const (
	imagesFieldSpecs = `
images:
- kind: Pod
  path: spec/initContainers
- kind: Pod
  path: spec/containers
- kind: Deployment
  path: spec/template/spec/initContainers
- kind: Deployment
  path: spec/template/spec/containers
`
)
<commit_msg>Leave defautconfig empty for images<commit_after>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package defaultconfig

const (
	// imageFieldSpecs is left empty since `containers` and `initContainers`
	// of *ANY* kind in *ANY* path are builtin supported in code
	imagesFieldSpecs = ``
)
<|endoftext|>