<commit_before>package main

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strconv"
	"strings"

	"github.com/nsf/termbox-go"
)

const (
	STACK_LINE = 0
	BODY_START = 2
)

func runServer() {

	l, err := net.Listen("unix", "/tmp/jeb-socket")
	if err != nil {
		log.Println(err)
		return
	}
	defer l.Close()

	fmt.Println("Waiting for message from client...")

	c, err := l.Accept()
	if err != nil {
		log.Println(err)
		return
	}
	defer c.Close()

	j := NewJeb(c, keyLoop())
	j.run() // doesn't return until exit
}

func setupLogging() {
	log.Println("Logging to jeb.log")
	l, err := os.Create("jeb.log")
	if err != nil {
		log.Println(err)
		return
	}
	log.SetOutput(l)
}

// Runs a go-routine to fetch keyboard events
func keyLoop() <-chan termbox.Event {
	ch := make(chan termbox.Event)
	go func() {
		var ev termbox.Event
		for {
			ev = termbox.PollEvent()
			if ev.Type == termbox.EventKey {
				ch <- ev
			}
		}
	}()
	return ch
}

type Jeb struct {
	c         net.Conn
	buf       []byte
	files     map[string][]string
	stack     *stack
	contLevel int
	keyChan   <-chan termbox.Event
}

func NewJeb(c net.Conn, keyChan <-chan termbox.Event) *Jeb {
	return &Jeb{
		c:         c,
		buf:       make([]byte, 1024),
		files:     make(map[string][]string),
		stack:     newStack(),
		contLevel: -1,
		keyChan:   keyChan,
	}
}

func (j *Jeb) run() {
	termbox.Init()
	defer termbox.Close()
	setupLogging()

	for {
		numRead := j.receive()
		if numRead == 0 {
			break
		}

		filename, line, function := j.parse(numRead)
		if filename == "" {
			j.proceed()
			continue
		}

		j.stack.condPush(function)

		if j.contLevel == -1 || j.contLevel < j.stack.Pos {
			j.display(filename, line)

			_, err := j.waitForInput()
			if err != nil {
				break
			}
			/*
				if input == 's' {
					// step in
					j.contLevel = -1
				} else if input == 'n' {
					// next (step over)
					j.contLevel = j.stack.Pos
				}
			*/

		}
		j.proceed()
	}
}

func (j *Jeb) receive() int {
	numRead, err := j.c.Read(j.buf)
	if err == io.EOF {
		return 0
	}
	if err != nil {
		log.Println(err)
		return -1
	}
	return numRead
}

// TODO: This should be able to return error as well
func (j *Jeb) parse(numRead int) (filename string, line int, function string) {
	var err error

	parts := strings.Split(string(j.buf[:numRead]), ":")
	cmd := parts[0]
	args := parts[1:]

	switch cmd {
	case "ENTER":
		j.stack.push(args[0])
		return
	case "EXIT":
		j.stack.pop(args[0])
		return
	case "LINE":
	default:
		log.Println("Unknown command: %s\n", cmd)
	}

	if len(args) != 3 {
		log.Println("Expected 3 parts, got %v", args)
	}
	filename = args[0]
	line, err = strconv.Atoi(args[1])
	if err != nil {
		log.Println(err)
	}
	function = args[2]

	return filename, line, function
}

// proceed tells client to continue to next line
func (j *Jeb) proceed() {
	j.c.Write([]byte{'\n'})
}

func (j *Jeb) waitForInput() (rune, error) {
	ev := <-j.keyChan
	if ev.Key == termbox.KeyCtrlC {
		return 0, errors.New("Ctrl-C")
	}
	return ev.Ch, nil
}

func (j *Jeb) display(filename string, lineNum int) {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	j.printStackLine()

	lines, ok := j.files[filename]
	if !ok {
		j.load(filename)
		lines, _ = j.files[filename]
	}

	_, h := termbox.Size()
	middle := h / 2
	lineStart := 0
	highlight := lineNum
	if lineNum > middle {
		lineStart = lineNum - middle
		highlight = middle
	}
	lineEnd := lineStart + h
	if lineEnd > len(lines) {
		lineEnd = len(lines)
	}

	for i, line := range lines[lineStart:lineEnd] {
		offset := i + BODY_START
		if i+1 == highlight {
			tbprint(0, offset, termbox.ColorDefault|termbox.AttrReverse, termbox.ColorDefault, line)
		} else {
			out(0, offset, line)
		}
	}
	termbox.Flush()
}

// load filename as array of file lines into Jeb.files
func (j *Jeb) load(filename string) {
	filedata, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Println(err)
		return
	}
	lines := strings.Split(string(filedata), "\n")

	j.files[filename] = lines
}

func (j *Jeb) printStackLine() {
	st := j.stack.stack()
	stackLine := make([]string, 0, len(st))
	for _, function := range st {
		stackLine = append(stackLine, function)
	}
	w, _ := termbox.Size()
	out(0, STACK_LINE, strings.Join(stackLine, "> "))
	out(0, STACK_LINE+1, strings.Repeat("-", w))
}

func bold(str string) string {
	return "\033[1m" + str + "\033[0m"
}

func out(x, y int, msg string) {
	tbprint(x, y, termbox.ColorDefault, termbox.ColorDefault, msg)
}
func tbprint(x, y int, fg, bg termbox.Attribute, msg string) {
	for _, c := range msg {
		termbox.SetCell(x, y, c, fg, bg)
		x++
	}
}

type stack struct {
	s   []string
	Pos int
}

func newStack() *stack {
	return &stack{
		s:   make([]string, 64),
		Pos: 0,
	}
}

func (s *stack) push(str string) {
	s.s[s.Pos] = str
	s.Pos++
}

func (s *stack) pop(str string) {
	s.Pos--
}

func (s *stack) stack() []string {
	return s.s[:s.Pos]
}

// condPush pushes 'str' on to the stack if it's not already the last element
func (s *stack) condPush(str string) {
	if s.Pos == 0 || s.s[s.Pos-1] != str {
		s.push(str)
	}
}
<commit_msg>Fix Println / Printf<commit_after>package main

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strconv"
	"strings"

	"github.com/nsf/termbox-go"
)

const (
	STACK_LINE = 0
	BODY_START = 2
)

func runServer() {

	l, err := net.Listen("unix", "/tmp/jeb-socket")
	if err != nil {
		log.Println(err)
		return
	}
	defer l.Close()

	fmt.Println("Waiting for message from client...")

	c, err := l.Accept()
	if err != nil {
		log.Println(err)
		return
	}
	defer c.Close()

	j := NewJeb(c, keyLoop())
	j.run() // doesn't return until exit
}

func setupLogging() {
	log.Println("Logging to jeb.log")
	l, err := os.Create("jeb.log")
	if err != nil {
		log.Println(err)
		return
	}
	log.SetOutput(l)
}

// Runs a go-routine to fetch keyboard events
func keyLoop() <-chan termbox.Event {
	ch := make(chan termbox.Event)
	go func() {
		var ev termbox.Event
		for {
			ev = termbox.PollEvent()
			if ev.Type == termbox.EventKey {
				ch <- ev
			}
		}
	}()
	return ch
}

type Jeb struct {
	c         net.Conn
	buf       []byte
	files     map[string][]string
	stack     *stack
	contLevel int
	keyChan   <-chan termbox.Event
}

func NewJeb(c net.Conn, keyChan <-chan termbox.Event) *Jeb {
	return &Jeb{
		c:         c,
		buf:       make([]byte, 1024),
		files:     make(map[string][]string),
		stack:     newStack(),
		contLevel: -1,
		keyChan:   keyChan,
	}
}

func (j *Jeb) run() {
	termbox.Init()
	defer termbox.Close()
	setupLogging()

	for {
		numRead := j.receive()
		if numRead == 0 {
			break
		}

		filename, line, function := j.parse(numRead)
		if filename == "" {
			j.proceed()
			continue
		}

		j.stack.condPush(function)

		if j.contLevel == -1 || j.contLevel < j.stack.Pos {
			j.display(filename, line)

			_, err := j.waitForInput()
			if err != nil {
				break
			}
			/*
				if input == 's' {
					// step in
					j.contLevel = -1
				} else if input == 'n' {
					// next (step over)
					j.contLevel = j.stack.Pos
				}
			*/

		}
		j.proceed()
	}
}

func (j *Jeb) receive() int {
	numRead, err := j.c.Read(j.buf)
	if err == io.EOF {
		return 0
	}
	if err != nil {
		log.Println(err)
		return -1
	}
	return numRead
}

// TODO: This should be able to return error as well
func (j *Jeb) parse(numRead int) (filename string, line int, function string) {
	var err error

	parts := strings.Split(string(j.buf[:numRead]), ":")
	cmd := parts[0]
	args := parts[1:]

	switch cmd {
	case "ENTER":
		j.stack.push(args[0])
		return
	case "EXIT":
		j.stack.pop(args[0])
		return
	case "LINE":
	default:
		log.Println("Unknown command:", cmd)
	}

	if len(args) != 3 {
		log.Println("Expected 3 parts, got", args)
	}
	filename = args[0]
	line, err = strconv.Atoi(args[1])
	if err != nil {
		log.Println(err)
	}
	function = args[2]

	return filename, line, function
}

// proceed tells client to continue to next line
func (j *Jeb) proceed() {
	j.c.Write([]byte{'\n'})
}

func (j *Jeb) waitForInput() (rune, error) {
	ev := <-j.keyChan
	if ev.Key == termbox.KeyCtrlC {
		return 0, errors.New("Ctrl-C")
	}
	return ev.Ch, nil
}

func (j *Jeb) display(filename string, lineNum int) {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	j.printStackLine()

	lines, ok := j.files[filename]
	if !ok {
		j.load(filename)
		lines, _ = j.files[filename]
	}

	_, h := termbox.Size()
	middle := h / 2
	lineStart := 0
	highlight := lineNum
	if lineNum > middle {
		lineStart = lineNum - middle
		highlight = middle
	}
	lineEnd := lineStart + h
	if lineEnd > len(lines) {
		lineEnd = len(lines)
	}

	for i, line := range lines[lineStart:lineEnd] {
		offset := i + BODY_START
		if i+1 == highlight {
			tbprint(0, offset, termbox.ColorDefault|termbox.AttrReverse, termbox.ColorDefault, line)
		} else {
			out(0, offset, line)
		}
	}
	termbox.Flush()
}

// load filename as array of file lines into Jeb.files
func (j *Jeb) load(filename string) {
	filedata, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Println(err)
		return
	}
	lines := strings.Split(string(filedata), "\n")

	j.files[filename] = lines
}

func (j *Jeb) printStackLine() {
	st := j.stack.stack()
	stackLine := make([]string, 0, len(st))
	for _, function := range st {
		stackLine = append(stackLine, function)
	}
	w, _ := termbox.Size()
	out(0, STACK_LINE, strings.Join(stackLine, "> "))
	out(0, STACK_LINE+1, strings.Repeat("-", w))
}

func bold(str string) string {
	return "\033[1m" + str + "\033[0m"
}

func out(x, y int, msg string) {
	tbprint(x, y, termbox.ColorDefault, termbox.ColorDefault, msg)
}
func tbprint(x, y int, fg, bg termbox.Attribute, msg string) {
	for _, c := range msg {
		termbox.SetCell(x, y, c, fg, bg)
		x++
	}
}

type stack struct {
	s   []string
	Pos int
}

func newStack() *stack {
	return &stack{
		s:   make([]string, 64),
		Pos: 0,
	}
}

func (s *stack) push(str string) {
	s.s[s.Pos] = str
	s.Pos++
}

func (s *stack) pop(str string) {
	s.Pos--
}

func (s *stack) stack() []string {
	return s.s[:s.Pos]
}

// condPush pushes 'str' on to the stack if it's not already the last element
func (s *stack) condPush(str string) {
	if s.Pos == 0 || s.s[s.Pos-1] != str {
		s.push(str)
	}
}
<|endoftext|>


<commit_before>package tarantool

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net"
	"sync"
)

const greetingSize = 128
const saltSize = 32
const tarantoolVersion = "Tarantool 1.6.8 (Binary)"
const connBufSize = 128 * 1024

type QueryHandler func(query Query) *Result
type OnDisconnectCallback func()

type IprotoServer struct {
	conn          net.Conn
	reader        *bufio.Reader
	writer        *bufio.Writer
	uuid          string
	salt          []byte // base64-encoded salt
	quit          chan bool
	handler       QueryHandler
	on_disconnect OnDisconnectCallback
	output        chan []byte
	closeOnce     sync.Once
}

func NewIprotoServer(uuid string, handler QueryHandler, on_disconnect OnDisconnectCallback) *IprotoServer {
	return &IprotoServer{
		conn:          nil,
		reader:        nil,
		writer:        nil,
		handler:       handler,
		on_disconnect: on_disconnect,
		uuid:          uuid,
	}
}

func (s *IprotoServer) Accept(conn net.Conn) {
	s.conn = conn
	s.reader = bufio.NewReaderSize(conn, connBufSize)
	s.writer = bufio.NewWriterSize(conn, connBufSize)
	s.quit = make(chan bool)
	s.output = make(chan []byte, 1024)

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

	go s.loop()
}

func (s *IprotoServer) CheckAuth(hash []byte, password string) bool {
	scr, err := scramble(s.salt, password)
	if err != nil {
		return false
	}

	if len(scr) != len(hash) {
		return false
	}

	for i, v := range hash {
		if v != scr[i] {
			return false
		}
	}
	return true
}

func (s *IprotoServer) Close() {
	s.closeOnce.Do(func() {
		if s.on_disconnect != nil {
			s.on_disconnect()
		}
		close(s.quit)
		s.conn.Close()
	})
}

func (s *IprotoServer) greet() (err error) {
	var line1, line2 string
	var format, greeting string
	var n int

	salt := make([]byte, saltSize)
	_, err = rand.Read(s.salt)
	if err != nil {
		return
	}

	s.salt = []byte(base64.StdEncoding.EncodeToString(salt))

	line1 = fmt.Sprintf("%s %s", tarantoolVersion, s.uuid)
	line2 = fmt.Sprintf("%s", s.salt)

	format = fmt.Sprintf("%%-%ds\n%%-%ds\n", greetingSize/2-1, greetingSize/2-1)
	greeting = fmt.Sprintf(format, line1, line2)

	// send greeting
	n, err = fmt.Fprintf(s.writer, "%s", greeting)
	if err != nil || n != greetingSize {
		return
	}

	return s.writer.Flush()
}

func (s *IprotoServer) loop() {
	go s.read()
	go s.write()
}

func (s *IprotoServer) read() {
	var packet *Packet
	var err error
	var body []byte

	r := s.reader

READER_LOOP:
	for {
		select {
		case <-s.quit:
			break READER_LOOP
		default:
			// read raw bytes
			body, err = readMessage(r)
			if err != nil {
				break READER_LOOP
			}

			packet, err = decodePacket(bytes.NewBuffer(body))
			if err != nil {
				break READER_LOOP
			}

			if packet.request != nil {
				go func(packet *Packet) {
					var res *Result
					var code = byte(packet.code)

					if code == PingRequest {
						s.output <- packIprotoOk(packet.requestID)
					} else {
						res = s.handler(packet.request.(Query))
						body, _ = res.pack(packet.requestID)
						s.output <- body
					}
				}(packet)
			}
		}
	}

	s.Close()
}

func (s *IprotoServer) write() {
	var err error
	var n int

	w := s.writer

WRITER_LOOP:
	for {
		select {
		case messageBody, ok := <-s.output:
			if !ok {
				break WRITER_LOOP
			}
			n, err = w.Write(messageBody)
			if err != nil || n != len(messageBody) {
				break WRITER_LOOP
			}
		case <-s.quit:
			w.Flush()
			break WRITER_LOOP
		default:
			if err = w.Flush(); err != nil {
				break WRITER_LOOP
			}

			// same without flush
			select {
			case messageBody, ok := <-s.output:
				if !ok {
					break WRITER_LOOP
				}
				n, err = w.Write(messageBody)
				if err != nil || n != len(messageBody) {
					break WRITER_LOOP
				}
			case <-s.quit:
				w.Flush()
				break WRITER_LOOP
			}

		}
	}

	s.Close()
}
<commit_msg>Rename OnDisconnect cb to OnClose<commit_after>package tarantool

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net"
	"sync"
)

const greetingSize = 128
const saltSize = 32
const tarantoolVersion = "Tarantool 1.6.8 (Binary)"
const connBufSize = 128 * 1024

type QueryHandler func(query Query) *Result
type OnCloseCallback func()

type IprotoServer struct {
	conn      net.Conn
	reader    *bufio.Reader
	writer    *bufio.Writer
	uuid      string
	salt      []byte // base64-encoded salt
	quit      chan bool
	handler   QueryHandler
	onClose   OnCloseCallback
	output    chan []byte
	closeOnce sync.Once
}

func NewIprotoServer(uuid string, handler QueryHandler, onClose OnCloseCallback) *IprotoServer {
	return &IprotoServer{
		conn:    nil,
		reader:  nil,
		writer:  nil,
		handler: handler,
		onClose: onClose,
		uuid:    uuid,
	}
}

func (s *IprotoServer) Accept(conn net.Conn) {
	s.conn = conn
	s.reader = bufio.NewReaderSize(conn, connBufSize)
	s.writer = bufio.NewWriterSize(conn, connBufSize)
	s.quit = make(chan bool)
	s.output = make(chan []byte, 1024)

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

	go s.loop()
}

func (s *IprotoServer) CheckAuth(hash []byte, password string) bool {
	scr, err := scramble(s.salt, password)
	if err != nil {
		return false
	}

	if len(scr) != len(hash) {
		return false
	}

	for i, v := range hash {
		if v != scr[i] {
			return false
		}
	}
	return true
}

func (s *IprotoServer) Close() {
	s.closeOnce.Do(func() {
		if s.onClose != nil {
			s.onClose()
		}
		close(s.quit)
		s.conn.Close()
	})
}

func (s *IprotoServer) greet() (err error) {
	var line1, line2 string
	var format, greeting string
	var n int

	salt := make([]byte, saltSize)
	_, err = rand.Read(s.salt)
	if err != nil {
		return
	}

	s.salt = []byte(base64.StdEncoding.EncodeToString(salt))

	line1 = fmt.Sprintf("%s %s", tarantoolVersion, s.uuid)
	line2 = fmt.Sprintf("%s", s.salt)

	format = fmt.Sprintf("%%-%ds\n%%-%ds\n", greetingSize/2-1, greetingSize/2-1)
	greeting = fmt.Sprintf(format, line1, line2)

	// send greeting
	n, err = fmt.Fprintf(s.writer, "%s", greeting)
	if err != nil || n != greetingSize {
		return
	}

	return s.writer.Flush()
}

func (s *IprotoServer) loop() {
	go s.read()
	go s.write()
}

func (s *IprotoServer) read() {
	var packet *Packet
	var err error
	var body []byte

	r := s.reader

READER_LOOP:
	for {
		select {
		case <-s.quit:
			break READER_LOOP
		default:
			// read raw bytes
			body, err = readMessage(r)
			if err != nil {
				break READER_LOOP
			}

			packet, err = decodePacket(bytes.NewBuffer(body))
			if err != nil {
				break READER_LOOP
			}

			if packet.request != nil {
				go func(packet *Packet) {
					var res *Result
					var code = byte(packet.code)

					if code == PingRequest {
						s.output <- packIprotoOk(packet.requestID)
					} else {
						res = s.handler(packet.request.(Query))
						body, _ = res.pack(packet.requestID)
						s.output <- body
					}
				}(packet)
			}
		}
	}

	s.Close()
}

func (s *IprotoServer) write() {
	var err error
	var n int

	w := s.writer

WRITER_LOOP:
	for {
		select {
		case messageBody, ok := <-s.output:
			if !ok {
				break WRITER_LOOP
			}
			n, err = w.Write(messageBody)
			if err != nil || n != len(messageBody) {
				break WRITER_LOOP
			}
		case <-s.quit:
			w.Flush()
			break WRITER_LOOP
		default:
			if err = w.Flush(); err != nil {
				break WRITER_LOOP
			}

			// same without flush
			select {
			case messageBody, ok := <-s.output:
				if !ok {
					break WRITER_LOOP
				}
				n, err = w.Write(messageBody)
				if err != nil || n != len(messageBody) {
					break WRITER_LOOP
				}
			case <-s.quit:
				w.Flush()
				break WRITER_LOOP
			}

		}
	}

	s.Close()
}
<|endoftext|>


<commit_before>package engineio

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"github.com/googollee/go-engine.io/polling"
	"github.com/googollee/go-engine.io/websocket"
	"net/http"
	"time"
)

type config struct {
	PingTimeout   time.Duration
	PingInterval  time.Duration
	AllowRequest  func(*http.Request) error
	AllowUpgrades bool
	Cookie        string
}

// Server is the server of engine.io.
type Server struct {
	config         config
	socketChan     chan Conn
	serverSessions *serverSessions
	creaters       transportCreaters
}

// NewServer returns the server suppported given transports. If transports is nil, server will use ["polling", "webosocket"] as default.
func NewServer(transports []string) (*Server, error) {
	if transports == nil {
		transports = []string{"polling", "websocket"}
	}
	creaters := make(transportCreaters)
	for _, t := range transports {
		switch t {
		case "polling":
			creaters[t] = polling.Creater
		case "websocket":
			creaters[t] = websocket.Creater
		default:
			return nil, InvalidError
		}
	}
	return &Server{
		config: config{
			PingTimeout:   60000 * time.Millisecond,
			PingInterval:  25000 * time.Millisecond,
			AllowRequest:  func(*http.Request) error { return nil },
			AllowUpgrades: true,
			Cookie:        "io",
		},
		socketChan:     make(chan Conn),
		serverSessions: newServerSessions(),
		creaters:       creaters,
	}, nil
}

// SetPingTimeout sets the timeout of ping. When time out, server will close connection. Default is 60s.
func (s *Server) SetPingTimeout(t time.Duration) {
	s.config.PingTimeout = t
}

// SetPingInterval sets the interval of ping. Default is 25s.
func (s *Server) SetPingInterval(t time.Duration) {
	s.config.PingInterval = t
}

// SetAllowRequest sets the middleware function when establish connection. If it return non-nil, connection won't be established. Default will allow all request.
func (s *Server) SetAllowRequest(f func(*http.Request) error) {
	s.config.AllowRequest = f
}

// SetAllowUpgrades sets whether server allows transport upgrade. Default is true.
func (s *Server) SetAllowUpgrades(allow bool) {
	s.config.AllowUpgrades = allow
}

// SetCookie sets the name of cookie which used by engine.io. Default is "io".
func (s *Server) SetCookie(prefix string) {
	s.config.Cookie = prefix
}

// ServeHTTP handles http request.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()

	cookies := r.Cookies()
	sid := r.URL.Query().Get("sid")
	conn := s.serverSessions.Get(sid)
	if conn == nil {
		if sid != "" {
			http.Error(w, "invalid sid", http.StatusBadRequest)
			return
		}

		if err := s.config.AllowRequest(r); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		sid = s.newId(r)
		var err error
		conn, err = newServerConn(sid, w, r, s)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		s.serverSessions.Set(sid, conn)
		cookies = append(cookies, &http.Cookie{
			Name:  s.config.Cookie,
			Value: sid,
		})

		s.socketChan <- conn
	}
	for _, c := range cookies {
		w.Header().Set("Set-Cookie", c.String())
	}
	conn.ServeHTTP(w, r)
}

// Accept returns Conn when client connect to server.
func (s *Server) Accept() (Conn, error) {
	return <-s.socketChan, nil
}

func (s *Server) configure() config {
	return s.config
}

func (s *Server) transports() transportCreaters {
	return s.creaters
}

func (s *Server) onClose(id string) {
	s.serverSessions.Remove(id)
}

func (s *Server) newId(r *http.Request) string {
	hash := fmt.Sprintf("%s %s", r.RemoteAddr, time.Now())
	buf := bytes.NewBuffer(nil)
	sum := md5.Sum([]byte(hash))
	encoder := base64.NewEncoder(base64.URLEncoding, buf)
	encoder.Write(sum[:])
	encoder.Close()
	return buf.String()[:20]
}
<commit_msg>added feature to set max connection<commit_after>package engineio

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"github.com/googollee/go-engine.io/polling"
	"github.com/googollee/go-engine.io/websocket"
	"net/http"
	"sync/atomic"
	"time"
)

type config struct {
	PingTimeout   time.Duration
	PingInterval  time.Duration
	MaxConnection int
	AllowRequest  func(*http.Request) error
	AllowUpgrades bool
	Cookie        string
}

// Server is the server of engine.io.
type Server struct {
	config         config
	socketChan     chan Conn
	serverSessions *serverSessions
	creaters       transportCreaters
}

// NewServer returns the server suppported given transports. If transports is nil, server will use ["polling", "webosocket"] as default.
func NewServer(transports []string) (*Server, error) {
	if transports == nil {
		transports = []string{"polling", "websocket"}
	}
	creaters := make(transportCreaters)
	for _, t := range transports {
		switch t {
		case "polling":
			creaters[t] = polling.Creater
		case "websocket":
			creaters[t] = websocket.Creater
		default:
			return nil, InvalidError
		}
	}
	return &Server{
		config: config{
			PingTimeout:   60000 * time.Millisecond,
			PingInterval:  25000 * time.Millisecond,
			MaxConnection: 1000,
			AllowRequest:  func(*http.Request) error { return nil },
			AllowUpgrades: true,
			Cookie:        "io",
		},
		socketChan:     make(chan Conn),
		serverSessions: newServerSessions(),
		creaters:       creaters,
	}, nil
}

// SetPingTimeout sets the timeout of ping. When time out, server will close connection. Default is 60s.
func (s *Server) SetPingTimeout(t time.Duration) {
	s.config.PingTimeout = t
}

// SetPingInterval sets the interval of ping. Default is 25s.
func (s *Server) SetPingInterval(t time.Duration) {
	s.config.PingInterval = t
}

// SetMaxConnection sets the max connetion. Default is 1000.
func (s *Server) SetMaxConnection(n int) {
	s.config.MaxConnection = n
}

// SetAllowRequest sets the middleware function when establish connection. If it return non-nil, connection won't be established. Default will allow all request.
func (s *Server) SetAllowRequest(f func(*http.Request) error) {
	s.config.AllowRequest = f
}

// SetAllowUpgrades sets whether server allows transport upgrade. Default is true.
func (s *Server) SetAllowUpgrades(allow bool) {
	s.config.AllowUpgrades = allow
}

// SetCookie sets the name of cookie which used by engine.io. Default is "io".
func (s *Server) SetCookie(prefix string) {
	s.config.Cookie = prefix
}

// ServeHTTP handles http request.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()

	cookies := r.Cookies()
	sid := r.URL.Query().Get("sid")
	conn := s.serverSessions.Get(sid)
	if conn == nil {
		if sid != "" {
			http.Error(w, "invalid sid", http.StatusBadRequest)
			return
		}

		if err := s.config.AllowRequest(r); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		n := atomic.AddInt32(&s.currentConnection, 1)
		if int(n) > s.config.MaxConnection {
			http.Error(w, "too many connections", http.StatusServiceUnavailable)
		}

		sid = s.newId(r)
		var err error
		conn, err = newServerConn(sid, w, r, s)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		s.serverSessions.Set(sid, conn)
		cookies = append(cookies, &http.Cookie{
			Name:  s.config.Cookie,
			Value: sid,
		})

		s.socketChan <- conn
	}
	for _, c := range cookies {
		w.Header().Set("Set-Cookie", c.String())
	}
	conn.ServeHTTP(w, r)
}

// Accept returns Conn when client connect to server.
func (s *Server) Accept() (Conn, error) {
	return <-s.socketChan, nil
}

func (s *Server) configure() config {
	return s.config
}

func (s *Server) transports() transportCreaters {
	return s.creaters
}

func (s *Server) onClose(id string) {
	s.serverSessions.Remove(id)
	atomic.AddInt32(&s.currentConnection, -1)
}

func (s *Server) newId(r *http.Request) string {
	hash := fmt.Sprintf("%s %s", r.RemoteAddr, time.Now())
	buf := bytes.NewBuffer(nil)
	sum := md5.Sum([]byte(hash))
	encoder := base64.NewEncoder(base64.URLEncoding, buf)
	encoder.Write(sum[:])
	encoder.Close()
	return buf.String()[:20]
}
<|endoftext|>


<commit_before>package main

import (
	r "github.com/dancannon/gorethink"
	"github.com/gorilla/mux"
	"log"
	"net/http"
	"time"
)

var (
	router  *mux.Router
	session *r.Session
)

func init() {
	var err error

	session, err = r.Connect(map[string]interface{}{
		"address":  "localhost:28015",
		"database": "todo",
	})
	if err != nil {
		log.Fatalln(err.Error())
	}
}

func NewServer(addr string) *http.Server {
	// Setup router
	router = initRouting()

	// Create and start server
	return &http.Server{
		Addr:    addr,
		Handler: router,
	}
}

func StartServer(server *http.Server) {
	err := server.ListenAndServe()
	if err != nil {
		log.Fatalln("Error: %v", err)
	}
}

func initRouting() *mux.Router {
	r := mux.NewRouter()

	r.HandleFunc("/", indexHandler)
	r.HandleFunc("/active", activeIndexHandler)
	r.HandleFunc("/completed", completedIndexHandler)
	r.HandleFunc("/new", newHandler)
	r.HandleFunc("/toggle/{id}", toggleHandler)
	r.HandleFunc("/delete/{id}", deleteHandler)
	r.HandleFunc("/clear", clearHandler)

	// Add handler for static files
	r.PathPrefix("/").Handler(http.FileServer(http.Dir("static")))

	return r
}

// Handlers

func indexHandler(w http.ResponseWriter, req *http.Request) {
	items := []TodoItem{}

	// Fetch all the items from the database
	rows, err := r.Table("items").OrderBy(r.Asc("Created")).Run(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Scan each row into a TodoItem instance and then add this to the list
	for rows.Next() {
		var item TodoItem

		err := rows.Scan(&item)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		items = append(items, item)
	}

	renderTemplate(w, "index", items)
}

func activeIndexHandler(w http.ResponseWriter, req *http.Request) {
	items := []TodoItem{}

	// Fetch all the items from the database
	query := r.Table("items").Filter(r.Row.Field("Status").Eq("active"))
	query = query.OrderBy(r.Asc("Created"))
	rows, err := query.Run(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Scan each row into a TodoItem instance and then add this to the list
	for rows.Next() {
		var item TodoItem

		err := rows.Scan(&item)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		items = append(items, item)
	}

	renderTemplate(w, "index", items)
}

func completedIndexHandler(w http.ResponseWriter, req *http.Request) {
	items := []TodoItem{}

	// Fetch all the items from the database
	query := r.Table("items").Filter(r.Row.Field("Status").Eq("complete"))
	query = query.OrderBy(r.Asc("Created"))
	rows, err := query.Run(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Scan each row into a TodoItem instance and then add this to the list
	for rows.Next() {
		var item TodoItem

		err := rows.Scan(&item)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		items = append(items, item)
	}

	renderTemplate(w, "index", items)
}

func newHandler(w http.ResponseWriter, req *http.Request) {
	// Create the item
	item := NewTodoItem(req.PostFormValue("text"))
	item.Created = time.Now()

	// Insert the new item into the database
	_, err := r.Table("items").Insert(item).RunWrite(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, req, "/", http.StatusFound)
}

func toggleHandler(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	if id == "" {
		http.NotFound(w, req)
		return
	}

	// Check that the item exists
	res, err := r.Table("items").Get(id).RunRow(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if res.IsNil() {
		http.NotFound(w, req)
		return
	}

	// Toggle the item
	_, err = r.Table("items").Get(id).Update(map[string]interface{}{"Status": r.Branch(
		r.Row.Field("Status").Eq("active"),
		"complete",
		"active",
	)}).RunWrite(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, req, "/", http.StatusFound)
}

func deleteHandler(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	if id == "" {
		http.NotFound(w, req)
		return
	}

	// Check that the item exists
	res, err := r.Table("items").Get(id).RunRow(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if res.IsNil() {
		http.NotFound(w, req)
		return
	}

	// Delete the item
	_, err = r.Table("items").Get(id).Delete().RunWrite(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, req, "/", http.StatusFound)
}

func clearHandler(w http.ResponseWriter, req *http.Request) {
	// Delete all completed items
	_, err := r.Table("items").Filter(
		r.Row.Field("Status").Eq("complete"),
	).Delete().RunWrite(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, req, "/", http.StatusFound)
}
<commit_msg>Updated to work with gorethink v0.3<commit_after>package main

import (
	r "github.com/dancannon/gorethink"
	"github.com/gorilla/mux"
	"log"
	"net/http"
	"time"
)

var (
	router  *mux.Router
	session *r.Session
)

func init() {
	var err error

	session, err = r.Connect(r.ConnectOpts{
		Address:  "localhost:28015",
		Database: "todo",
	})
	if err != nil {
		log.Fatalln(err.Error())
	}
}

func NewServer(addr string) *http.Server {
	// Setup router
	router = initRouting()

	// Create and start server
	return &http.Server{
		Addr:    addr,
		Handler: router,
	}
}

func StartServer(server *http.Server) {
	err := server.ListenAndServe()
	if err != nil {
		log.Fatalln("Error: %v", err)
	}
}

func initRouting() *mux.Router {
	r := mux.NewRouter()

	r.HandleFunc("/", indexHandler)
	r.HandleFunc("/active", activeIndexHandler)
	r.HandleFunc("/completed", completedIndexHandler)
	r.HandleFunc("/new", newHandler)
	r.HandleFunc("/toggle/{id}", toggleHandler)
	r.HandleFunc("/delete/{id}", deleteHandler)
	r.HandleFunc("/clear", clearHandler)

	// Add handler for static files
	r.PathPrefix("/").Handler(http.FileServer(http.Dir("static")))

	return r
}

// Handlers

func indexHandler(w http.ResponseWriter, req *http.Request) {
	items := []TodoItem{}

	// Fetch all the items from the database
	res, err := r.Table("items").OrderBy(r.Asc("Created")).Run(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	err = res.All(&items)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	renderTemplate(w, "index", items)
}

func activeIndexHandler(w http.ResponseWriter, req *http.Request) {
	items := []TodoItem{}

	// Fetch all the items from the database
	query := r.Table("items").Filter(r.Row.Field("Status").Eq("active"))
	query = query.OrderBy(r.Asc("Created"))
	res, err := query.Run(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	err = res.All(&items)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	renderTemplate(w, "index", items)
}

func completedIndexHandler(w http.ResponseWriter, req *http.Request) {
	items := []TodoItem{}

	// Fetch all the items from the database
	query := r.Table("items").Filter(r.Row.Field("Status").Eq("complete"))
	query = query.OrderBy(r.Asc("Created"))
	res, err := query.Run(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	err = res.All(&items)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	renderTemplate(w, "index", items)
}

func newHandler(w http.ResponseWriter, req *http.Request) {
	// Create the item
	item := NewTodoItem(req.PostFormValue("text"))
	item.Created = time.Now()

	// Insert the new item into the database
	_, err := r.Table("items").Insert(item).RunWrite(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, req, "/", http.StatusFound)
}

func toggleHandler(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	if id == "" {
		http.NotFound(w, req)
		return
	}

	// Check that the item exists
	res, err := r.Table("items").Get(id).Run(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if res.IsNil() {
		http.NotFound(w, req)
		return
	}

	// Toggle the item
	_, err = r.Table("items").Get(id).Update(map[string]interface{}{"Status": r.Branch(
		r.Row.Field("Status").Eq("active"),
		"complete",
		"active",
	)}).RunWrite(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, req, "/", http.StatusFound)
}

func deleteHandler(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	if id == "" {
		http.NotFound(w, req)
		return
	}

	// Check that the item exists
	res, err := r.Table("items").Get(id).Run(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if res.IsNil() {
		http.NotFound(w, req)
		return
	}

	// Delete the item
	_, err = r.Table("items").Get(id).Delete().RunWrite(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, req, "/", http.StatusFound)
}

func clearHandler(w http.ResponseWriter, req *http.Request) {
	// Delete all completed items
	_, err := r.Table("items").Filter(
		r.Row.Field("Status").Eq("complete"),
	).Delete().RunWrite(session)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, req, "/", http.StatusFound)
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"

	"github.com/golang/glog"
	"github.com/gorilla/context"

	"github.com/prashanthrv/sangeeblog/controllers"
	"github.com/prashanthrv/sangeeblog/system"

	"github.com/zenazn/goji"
	"github.com/zenazn/goji/graceful"
	"github.com/zenazn/goji/web"
	//"golang.org/x/net/http2"
)

func main() {
	// filename := flag.String("config", "config.toml", "Path to configuration file")
	//
	// flag.Parse()
	filename := "config.toml"
	defer glog.Flush()

	var application = &system.Application{}

	application.Init(&filename)
	application.LoadTemplates()

	// Setup static files
	static := web.New()
	publicPath := application.Config.Get("general.public_path").(string)
	static.Get("/assets/*", http.StripPrefix("/assets/", http.FileServer(http.Dir(publicPath))))

	http.Handle("/assets/", static)

	// Apply middleware
	goji.Use(application.ApplyTemplates)
	goji.Use(application.ApplySessions)
	//goji.Use(application.ApplyDbMap)
	goji.Use(application.ApplyGormDB)
	//goji.Use(application.ApplyAuth)
	goji.Use(application.ApplyIsXhr)
	goji.Use(application.ApplyCsrfProtection)
	goji.Use(context.ClearHandler)

	controller := &controllers.MainController{}

	// Couple of files - in the real world you would use nginx to serve them.
	goji.Get("/robots.txt", http.FileServer(http.Dir(publicPath)))
	goji.Get("/favicon.ico", http.FileServer(http.Dir(publicPath+"/images")))

	// Home page
	goji.Get("/", application.Route(controller, "Blog"))

	// Sign In routes
	goji.Get("/signin", application.Route(controller, "SignIn"))
	goji.Post("/signin", application.Route(controller, "SignInPost"))

	// Sign Up routes
	goji.Get("/signup", application.Route(controller, "SignUp"))
	goji.Post("/signup", application.Route(controller, "SignUpPost"))

	// KTHXBYE
	goji.Get("/logout", application.Route(controller, "Logout"))

	goji.Get("/blog", application.Route(controller, "Blog"))
	goji.Get("/post/:postid", application.Route(controller, "Post"))
	goji.Get("/category/:categoryid", application.Route(controller, "Categories"))
	goji.Get("/page/:pageid", application.Route(controller, "Pages"))

	graceful.PostHook(func() {
		application.Close()
	})
	goji.Serve()
	//http2.ConfigureServer(http, &http2.Server{})
}
<commit_msg>Middleware issue<commit_after>package main

import (
	"net/http"

	"github.com/golang/glog"
	"github.com/gorilla/context"

	"github.com/prashanthrv/sangeeblog/controllers"
	"github.com/prashanthrv/sangeeblog/system"

	"github.com/zenazn/goji"
	"github.com/zenazn/goji/graceful"
	"github.com/zenazn/goji/web"
	//"golang.org/x/net/http2"
)

func main() {
	// filename := flag.String("config", "config.toml", "Path to configuration file")
	//
	// flag.Parse()
	filename := "config.toml"
	defer glog.Flush()

	var application = &system.Application{}

	application.Init(&filename)
	application.LoadTemplates()

	// Setup static files
	static := web.New()
	publicPath := application.Config.Get("general.public_path").(string)
	static.Get("/assets/*", http.StripPrefix("/assets/", http.FileServer(http.Dir(publicPath))))

	http.Handle("/assets/", static)

	// Apply middleware
	//goji.Use(application.ApplyTemplates)
	//goji.Use(application.ApplySessions)
	//goji.Use(application.ApplyDbMap)
	goji.Use(application.ApplyGormDB)
	//goji.Use(application.ApplyAuth)
	//goji.Use(application.ApplyIsXhr)
	//goji.Use(application.ApplyCsrfProtection)
	//goji.Use(context.ClearHandler)

	controller := &controllers.MainController{}

	// Couple of files - in the real world you would use nginx to serve them.
	goji.Get("/robots.txt", http.FileServer(http.Dir(publicPath)))
	goji.Get("/favicon.ico", http.FileServer(http.Dir(publicPath+"/images")))

	// Home page
	goji.Get("/", application.Route(controller, "Blog"))

	// Sign In routes
	goji.Get("/signin", application.Route(controller, "SignIn"))
	goji.Post("/signin", application.Route(controller, "SignInPost"))

	// Sign Up routes
	goji.Get("/signup", application.Route(controller, "SignUp"))
	goji.Post("/signup", application.Route(controller, "SignUpPost"))

	// KTHXBYE
	goji.Get("/logout", application.Route(controller, "Logout"))

	goji.Get("/blog", application.Route(controller, "Blog"))
	goji.Get("/post/:postid", application.Route(controller, "Post"))
	goji.Get("/category/:categoryid", application.Route(controller, "Categories"))
	goji.Get("/page/:pageid", application.Route(controller, "Pages"))

	graceful.PostHook(func() {
		application.Close()
	})
	goji.Serve()
	//http2.ConfigureServer(http, &http2.Server{})
}
<|endoftext|>


<commit_before>package main

//go:generate go-bindata-assetfs static/...

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

type nameID struct {
	Name string
	ID   string
}

func sortedBookPropertyList(books []Book, getNameID func(Book) nameID, filterNameID func(nameID) bool, sortNameID func(nameID, nameID) bool) []nameID {
	doneItems := map[string]bool{}
	items := []nameID{}
	for _, b := range books {
		nid := getNameID(b)
		if doneItems[nid.ID] {
			continue
		}
		doneItems[nid.ID] = true
		items = append(items, nameID{
			Name: nid.Name,
			ID:   nid.ID,
		})
	}
	filteredItems := []nameID{}
	for _, ni := range items {
		if filterNameID(ni) {
			filteredItems = append(filteredItems, ni)
		}
	}
	sort.Slice(filteredItems, func(i, j int) bool {
		return sortNameID(filteredItems[i], filteredItems[j])
	})
	return filteredItems
}

func sortedBookList(books []Book, filterBook func(Book) bool, sortBook func(Book, Book) bool) []Book {
	filteredItems := []Book{}
	for _, book := range books {
		if filterBook(book) {
			filteredItems = append(filteredItems, book)
		}
	}
	sort.Slice(filteredItems, func(i, j int) bool {
		return sortBook(filteredItems[i], filteredItems[j])
	})
	return filteredItems
}

// DownloadHandler handles file download
func DownloadHandler(w http.ResponseWriter, r *http.Request) {
	bid := filepath.Base(r.URL.Path)
	bid = strings.Replace(bid, filepath.Ext(bid), "", 1)

	if bid == "download" {
		w.Header().Set("Content-Type", "text/html")
		for _, b := range books {
			io.WriteString(w, fmt.Sprintf("%s
", b.ID, b.Title))
		}
		return
	}

	for _, b := range books {
		if b.ID == bid {
			rd, err := os.Open(b.Filepath)
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				io.WriteString(w, "Error handling request")
				log.Printf("Error handling request for %s: %s\n", r.URL.Path, err)
			}

			w.Header().Set("Content-Disposition", "attachment; filename="+url.PathEscape(b.Title)+".epub")
			w.Header().Set("Content-Type", "application/epub+zip")
			_, err = io.Copy(w, rd)
			rd.Close()
			if err != nil {
				log.Printf("Error handling request for %s: %s\n", r.URL.Path, err)
			}
			return
		}
	}

	w.WriteHeader(http.StatusNotFound)
	io.WriteString(w, pageHTML("Not Found", "Could not find book with id "+bid))
}

// AuthorsHandler handles the authors page
func AuthorsHandler(w http.ResponseWriter, r *http.Request) {
	aid := filepath.Base(r.URL.Path)

	if aid == "authors" {
		w.Header().Set("Content-Type", "text/html")
		var listHTML bytes.Buffer

		authors := sortedBookPropertyList(books, func(b Book) nameID {
			return nameID{
				Name: b.Author,
				ID:   b.AuthorID,
			}
		}, func(ni nameID) bool {
			return ni.Name != ""
		}, func(a nameID, b nameID) bool {
			return a.Name < b.Name
		})
		for _, ni := range authors {
			listHTML.WriteString(itemCardHTML(ni.Name, "", "/authors/"+ni.ID))
		}

		io.WriteString(w, pageHTML("Authors", listHTML.String()))
		return
	}

	w.Header().Set("Content-Type", "text/html")

	matched := sortedBookList(books, func(book Book) bool {
		return book.AuthorID == aid
	}, func(a Book, b Book) bool {
		return a.Title < b.Title
	})

	aname := ""
	if len(matched) != 0 {
		aname = matched[0].Author
	}

	html, notfound := bookListPageHTML(matched, aname, "Author not found")

	if notfound {
		w.WriteHeader(http.StatusNotFound)
	}

	io.WriteString(w, html)
}

// SeriesHandler handles the series page
func SeriesHandler(w http.ResponseWriter, r *http.Request) {
	sid := filepath.Base(r.URL.Path)

	if sid == "series" {
		w.Header().Set("Content-Type", "text/html")
		var listHTML bytes.Buffer

		series := sortedBookPropertyList(books, func(b Book) nameID {
			return nameID{
				Name: b.Series.Name,
				ID:   b.Series.ID,
			}
		}, func(ni nameID) bool {
			return ni.Name != ""
		}, func(a nameID, b nameID) bool {
			return a.Name < b.Name
		})
		for _, ni := range series {
			listHTML.WriteString(itemCardHTML(ni.Name, "", "/series/"+ni.ID))
		}
		if len(series) == 0 {
			io.WriteString(w, pageHTML("Series", "No series have been found."))
			return
		}

		io.WriteString(w, pageHTML("Series", listHTML.String()))
		return
	}

	w.Header().Set("Content-Type", "text/html")

	matched := sortedBookList(books, func(book Book) bool {
		return book.Series.ID == sid
	}, func(a Book, b Book) bool {
		return a.Series.Index < b.Series.Index
	})

	sname := ""
	if len(matched) != 0 {
		sname = matched[0].Series.Name
	}

	html, notfound := bookListPageHTML(matched, sname, "Series not found")

	if notfound {
		w.WriteHeader(http.StatusNotFound)
	}

	io.WriteString(w, html)
}

// BooksHandler handles the books page
func BooksHandler(w http.ResponseWriter, r *http.Request) {
	bid := filepath.Base(r.URL.Path)

	if bid == "books" {
		w.Header().Set("Content-Type", "text/html")

		matched := sortedBookList(books, func(book Book) bool {
			return true
		}, func(a Book, b Book) bool {
			return a.ModTime.Unix() > b.ModTime.Unix()
		})

		html, notfound := bookListPageHTML(matched, "Books", "There are no books in your library.")

		if notfound {
			w.WriteHeader(http.StatusNotFound)
		}

		io.WriteString(w, html)
		return
	}

	w.Header().Set("Content-Type", "text/html")
	for _, b := range books {
		if b.ID == bid {
			io.WriteString(w, pageHTML(b.Title, bookHTML(&b, false)))
			return
		}
	}

	w.WriteHeader(http.StatusNotFound)
	io.WriteString(w, pageHTML("Not Found", "Could not find book with id "+bid))
}

// SearchHandler handles the search page
func SearchHandler(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query().Get("q")
	ql := strings.ToLower(q)

	if len(q) != 0 {
		w.Header().Set("Content-Type", "text/html")
		var booksHTML bytes.Buffer
		booksHTML.WriteString(`<form role="search" method="GET" action="/search/">

<input type="text" class="form-control" placeholder="Search" name="q" id="q" value="` + strings.Replace(q, `"`, """, -1) + `">

<button class="btn btn-default" type="submit"></button>

</form>
`)
		booksHTML.WriteString(`<div class="books cards">`)
		matched := false
		for _, b := range books {
			matches := false
			matches = matches || strings.Contains(strings.ToLower(b.Author), ql)
			matches = matches || strings.Contains(strings.ToLower(b.Title), ql)
			matches = matches || strings.Contains(strings.ToLower(b.Series.Name), ql)

			if matches {
				booksHTML.WriteString(bookHTML(&b, true))
				matched = true
			}
		}
		booksHTML.WriteString(`</div>`)
		if !matched {
			booksHTML.WriteString("No books matching your query have been found.")
			return
		}
		io.WriteString(w, pageHTML("Search Results: "+q, booksHTML.String()))
	} else {
		w.Header().Set("Content-Type", "text/html")
		io.WriteString(w, pageHTML("Search", `<form role="search" method="GET" action="/search/">

<input type="text" class="form-control" placeholder="Search" name="q" id="q">

<button class="btn btn-default" type="submit"></button>

</form>`))
	}
}

// JSONHandler handles the books.json file
func JSONHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	b, _ := json.Marshal(books)
	w.Write(b)
}

var books []Book

func runServer(bks []Book, addr string) {
	books = bks

	StaticHandler := http.StripPrefix("/static", http.FileServer(assetFS()))
	CoversHandler := http.StripPrefix("/covers", http.FileServer(http.Dir(*tempdir)))

	http.Handle("/static/", StaticHandler)
	http.Handle("/covers/", CoversHandler)
	http.HandleFunc("/download/", DownloadHandler)
	http.HandleFunc("/authors/", AuthorsHandler)
	http.HandleFunc("/series/", SeriesHandler)
	http.HandleFunc("/books/", BooksHandler)
	http.HandleFunc("/search/", SearchHandler)
	http.HandleFunc("/books.json", JSONHandler)
	http.HandleFunc("/random/", func(w http.ResponseWriter, r *http.Request) {
		rand.Seed(time.Now().Unix())
		n := rand.Int() % len(books)
		http.Redirect(w, r, "/books/"+books[n].ID, http.StatusTemporaryRedirect)
	})
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/books/", http.StatusTemporaryRedirect)
	})

	log.Printf("Serving on %s\n", addr)
	err := http.ListenAndServe(addr, nil)
	if err != nil {
		log.Fatalf("Error starting server: %s\n", err)
	}
}
<commit_msg>Improved download list<commit_after>package main

//go:generate go-bindata-assetfs static/...

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

type nameID struct {
	Name string
	ID   string
}

func sortedBookPropertyList(books []Book, getNameID func(Book) nameID, filterNameID func(nameID) bool, sortNameID func(nameID, nameID) bool) []nameID {
	doneItems := map[string]bool{}
	items := []nameID{}
	for _, b := range books {
		nid := getNameID(b)
		if doneItems[nid.ID] {
			continue
		}
		doneItems[nid.ID] = true
		items = append(items, nameID{
			Name: nid.Name,
			ID:   nid.ID,
		})
	}
	filteredItems := []nameID{}
	for _, ni := range items {
		if filterNameID(ni) {
			filteredItems = append(filteredItems, ni)
		}
	}
	sort.Slice(filteredItems, func(i, j int) bool {
		return sortNameID(filteredItems[i], filteredItems[j])
	})
	return filteredItems
}

func sortedBookList(books []Book, filterBook func(Book) bool, sortBook func(Book, Book) bool) []Book {
	filteredItems := []Book{}
	for _, book := range books {
		if filterBook(book) {
			filteredItems = append(filteredItems, book)
		}
	}
	sort.Slice(filteredItems, func(i, j int) bool {
		return sortBook(filteredItems[i], filteredItems[j])
	})
	return filteredItems
}

// DownloadHandler handles file download
func DownloadHandler(w http.ResponseWriter, r *http.Request) {
	bid := filepath.Base(r.URL.Path)
	bid = strings.Replace(bid, filepath.Ext(bid), "", 1)

	if bid == "download" {
		w.Header().Set("Content-Type", "text/html")
		sbl := sortedBookList(books, func(b Book) bool {
			return true
		}, func(a Book, b Book) bool {
			return a.Title < b.Title
		})
		for _, b := range sbl {
			io.WriteString(w, fmt.Sprintf("%s - %s - %s - %s
", b.ID, b.Title, b.Author, b.Series.Name, b.Series.Index))
		}
		return
	}

	for _, b := range books {
		if b.ID == bid {
			rd, err := os.Open(b.Filepath)
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				io.WriteString(w, "Error handling request")
				log.Printf("Error handling request for %s: %s\n", r.URL.Path, err)
			}

			w.Header().Set("Content-Disposition", "attachment; filename="+url.PathEscape(b.Title)+".epub")
			w.Header().Set("Content-Type", "application/epub+zip")
			_, err = io.Copy(w, rd)
			rd.Close()
			if err != nil {
				log.Printf("Error handling request for %s: %s\n", r.URL.Path, err)
			}
			return
		}
	}

	w.WriteHeader(http.StatusNotFound)
	io.WriteString(w, pageHTML("Not Found", "Could not find book with id "+bid))
}

// AuthorsHandler handles the authors page
func AuthorsHandler(w http.ResponseWriter, r *http.Request) {
	aid := filepath.Base(r.URL.Path)

	if aid == "authors" {
		w.Header().Set("Content-Type", "text/html")
		var listHTML bytes.Buffer

		authors := sortedBookPropertyList(books, func(b Book) nameID {
			return nameID{
				Name: b.Author,
				ID:   b.AuthorID,
			}
		}, func(ni nameID) bool {
			return ni.Name != ""
		}, func(a nameID, b nameID) bool {
			return a.Name < b.Name
		})
		for _, ni := range authors {
			listHTML.WriteString(itemCardHTML(ni.Name, "", "/authors/"+ni.ID))
		}

		io.WriteString(w, pageHTML("Authors", listHTML.String()))
		return
	}

	w.Header().Set("Content-Type", "text/html")

	matched := sortedBookList(books, func(book Book) bool {
		return book.AuthorID == aid
	}, func(a Book, b Book) bool {
		return a.Title < b.Title
	})

	aname := ""
	if len(matched) != 0 {
		aname = matched[0].Author
	}

	html, notfound := bookListPageHTML(matched, aname, "Author not found")

	if notfound {
		w.WriteHeader(http.StatusNotFound)
	}

	io.WriteString(w, html)
}

// SeriesHandler handles the series page
func SeriesHandler(w http.ResponseWriter, r *http.Request) {
	sid := filepath.Base(r.URL.Path)

	if sid == "series" {
		w.Header().Set("Content-Type", "text/html")
		var listHTML bytes.Buffer

		series := sortedBookPropertyList(books, func(b Book) nameID {
			return nameID{
				Name: b.Series.Name,
				ID:   b.Series.ID,
			}
		}, func(ni nameID) bool {
			return ni.Name != ""
		}, func(a nameID, b nameID) bool {
			return a.Name < b.Name
		})
		for _, ni := range series {
			listHTML.WriteString(itemCardHTML(ni.Name, "", "/series/"+ni.ID))
		}
		if len(series) == 0 {
			io.WriteString(w, pageHTML("Series", "No series have been found."))
			return
		}

		io.WriteString(w, pageHTML("Series", listHTML.String()))
		return
	}

	w.Header().Set("Content-Type", "text/html")

	matched := sortedBookList(books, func(book Book) bool {
		return book.Series.ID == sid
	}, func(a Book, b Book) bool {
		return a.Series.Index < b.Series.Index
	})

	sname := ""
	if len(matched) != 0 {
		sname = matched[0].Series.Name
	}

	html, notfound := bookListPageHTML(matched, sname, "Series not found")

	if notfound {
		w.WriteHeader(http.StatusNotFound)
	}

	io.WriteString(w, html)
}

// BooksHandler handles the books page
func BooksHandler(w http.ResponseWriter, r *http.Request) {
	bid := filepath.Base(r.URL.Path)

	if bid == "books" {
		w.Header().Set("Content-Type", "text/html")

		matched := sortedBookList(books, func(book Book) bool {
			return true
		}, func(a Book, b Book) bool {
			return a.ModTime.Unix() > b.ModTime.Unix()
		})

		html, notfound := bookListPageHTML(matched, "Books", "There are no books in your library.")

		if notfound {
			w.WriteHeader(http.StatusNotFound)
		}

		io.WriteString(w, html)
		return
	}

	w.Header().Set("Content-Type", "text/html")
	for _, b := range books {
		if b.ID == bid {
			io.WriteString(w, pageHTML(b.Title, bookHTML(&b, false)))
			return
		}
	}

	w.WriteHeader(http.StatusNotFound)
	io.WriteString(w, pageHTML("Not Found", "Could not find book with id "+bid))
}

// SearchHandler handles the search page
func SearchHandler(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query().Get("q")
	ql := strings.ToLower(q)

	if len(q) != 0 {
		w.Header().Set("Content-Type", "text/html")
		var booksHTML bytes.Buffer
		booksHTML.WriteString(`<form role="search" method="GET" action="/search/">

<input type="text" class="form-control" placeholder="Search" name="q" id="q" value="` + strings.Replace(q, `"`, """, -1) + `">

<button class="btn btn-default" type="submit"></button>

</form>
`)
		booksHTML.WriteString(`<div class="books cards">`)
		matched := false
		for _, b := range books {
			matches := false
			matches = matches || strings.Contains(strings.ToLower(b.Author), ql)
			matches = matches || strings.Contains(strings.ToLower(b.Title), ql)
			matches = matches || strings.Contains(strings.ToLower(b.Series.Name), ql)

			if matches {
				booksHTML.WriteString(bookHTML(&b, true))
				matched = true
			}
		}
		booksHTML.WriteString(`</div>`)
		if !matched {
			booksHTML.WriteString("No books matching your query have been found.")
			return
		}
		io.WriteString(w, pageHTML("Search Results: "+q, booksHTML.String()))
	} else {
		w.Header().Set("Content-Type", "text/html")
		io.WriteString(w, pageHTML("Search", `<form role="search" method="GET" action="/search/">

<input type="text" class="form-control" placeholder="Search" name="q" id="q">

<button class="btn btn-default" type="submit"></button>

</form>`))
	}
}

// JSONHandler handles the books.json file
func JSONHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	b, _ := json.Marshal(books)
	w.Write(b)
}

var books []Book

func runServer(bks []Book, addr string) {
	books = bks

	StaticHandler := http.StripPrefix("/static", http.FileServer(assetFS()))
	CoversHandler := http.StripPrefix("/covers", http.FileServer(http.Dir(*tempdir)))

	http.Handle("/static/", StaticHandler)
	http.Handle("/covers/", CoversHandler)
	http.HandleFunc("/download/", DownloadHandler)
	http.HandleFunc("/authors/", AuthorsHandler)
	http.HandleFunc("/series/", SeriesHandler)
	http.HandleFunc("/books/", BooksHandler)
	http.HandleFunc("/search/", SearchHandler)
	http.HandleFunc("/books.json", JSONHandler)
	http.HandleFunc("/random/", func(w http.ResponseWriter, r *http.Request) {
		rand.Seed(time.Now().Unix())
		n := rand.Int() % len(books)
		http.Redirect(w, r, "/books/"+books[n].ID, http.StatusTemporaryRedirect)
	})
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/books/", http.StatusTemporaryRedirect)
	})

	log.Printf("Serving on %s\n", addr)
	err := http.ListenAndServe(addr, nil)
	if err != nil {
		log.Fatalf("Error starting server: %s\n", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/mux"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

//curl -H "Content-Type: application/json" -XPUT -d '{"P"}' -f -v http://localhost:3000/beacon/{id}

const (
	dbName = "artroomServer"
	//DBURI for the mongodb
	DBURI = "127.0.0.1"

	waltersAPIPrefix    = "http://api.thewalters.org/v1/objects?apikey=ShxvahaBFNIcfWR7E78xXdssKIlXtUAJk9rDDrrmlvbOlxQKtASCzV4op5aHv2Il&title="
	waltersImagePrefix  = "http://static.thewalters.org/images/"
	waltersImagePostfix = "?width=500"
)

var (
	s = Server{}
	//CollectionNames is the shiz
	CollectionNames = []string{"beacon", "art"}
	beaconIndex     = mgo.Index{
		Key:        []string{"MinorID"},
		Unique:     true,
		DropDups:   true,
		Background: true,
		Sparse:     true,
		Name:       "beaconIndex",
		//ExpireAf
	}

	artIndex = mgo.Index{
		Key:        []string{"Beacon", "Title"},
		Unique:     true,
		DropDups:   true,
		Background: true,
		Sparse:     true,
		Name:       "ArtIndex",
		//ExpireAf
	}
	indices = []mgo.Index{beaconIndex, artIndex}
)

type (
	//Server is the name for the server deal wit it
	Server struct {
		Session *mgo.Session // The main session we'll we be cloning
		DBURI   string       // Where the DB is on the network
		dbName  string       // Name of the MongoDB
	}

	//Beacon is the struct that structures what the data for a beacon will look like
	Beacon struct {
		ProxID  string
		MajorID int
		MinorID int
	}
	//FormResponse is what we get back from the form
	FormResponse struct {
		Date        string
		Author      string
		Title       string
		ProxID      string
		MajorID     int
		MinorID     int
		Description string
	}

	//The structure the walter api gives us more or less
	WalterObj struct {
		ObjectID       int
		Collection     string
		Title          string
		Author         string
		Medium         string
		Description    string
		Images         string
		CuratorComment string
		Beacon         Beacon
		ImageURL       string
	}
	//Walters objects
	Walters struct {
		Items []WalterObj
	}
)

func main() {
	log.Println("Server is warming up...")
	initDB()
	defer s.Session.Close()
	http.Handle("/", initHandlers())
	log.Fatalln(http.ListenAndServe(":3000", nil))
}

func initHandlers() *mux.Router {
	r := mux.NewRouter()
	r.HandleFunc("/beacon", handlePOSTBeacon).Methods("POST", "OPTIONS")
	r.HandleFunc("/beacon/{minorID:[0-9]+}", handleBeacon).Methods("GET")
	//	cors := tigertonic.NewCORSBuilder().AddAllowedOrigins("*").AddAllowedHeaders("Access-Control-Allow-Headers", )
	//
	return r
}

func handlePOSTBeacon(w http.ResponseWriter, req *http.Request) {
	log.Println("I got here?")
	switch req.Method {
	case "OPTIONS":
		w.Header().Set("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		fallthrough
	case "POST":
		//		w.Header().Set("Access-Control-Allow-Origin", "*")
		var formResponse FormResponse
		err := ReadJSON(req, &formResponse)
		log.Printf("We have begun the ritual, %v", formResponse)
		client := &http.Client{}
		request, _ := http.NewRequest("GET", waltersAPIPrefix+strings.Replace(formResponse.Title, " ", "%20", -1), nil)
		request.Header.Add("Accept", "application/json")
		res, err := client.Do(request)
		if err != nil {
			http.Error(w, "Req didn't go through", 500)
			return
			//		log.Fatalf("Req didn't go through, %v", err)
		}
		defer res.Body.Close()
		var data json.RawMessage
		err = json.NewDecoder(res.Body).Decode(&data)
		if err != nil {
			http.Error(w, "Couldn't decode data...", 500)
			return
			//		log.Fatalf("Can't resolve the datum?%v", err)
		}
		wal := Walters{}
		err = json.Unmarshal(data, &wal)
		if err != nil {
			http.Error(w, "Couldn't marshall into individual waltOBJ.", 500)
			return
			//		log.Fatalf("Can't unmarshall into individ waltOBJ, %v", err)
		}
		if len(wal.Items) < 1 {
			http.Error(w, "Couldn't marshall correctly", 500)
			return
		}
		waltersObject := wal.Items[0]
		waltersObject.Beacon = Beacon{formResponse.ProxID, formResponse.MajorID, formResponse.MinorID}
		waltersObject.ImageURL = waltersImagePrefix + string(waltersObject.Images[0]) + waltersImagePostfix
		waltersObject.CuratorComment = formResponse.Description
		err = Insert("beacon", waltersObject.Beacon)
		if err != nil {
			//Return a 500
			http.Error(w, "Couldn't insert corresponding beacon.", 500)
			return
		}

		err = Insert("art", waltersObject)
		if err != nil {
			//Return a 500
			http.Error(w, "Couldn't insert corresponding art piece.", 500)
			return
		}
		http.Error(w, http.StatusText(http.StatusCreated), http.StatusCreated)
		break
	}
}

func handleBeacon(w http.ResponseWriter, req *http.Request) {
	params := mux.Vars(req)
	minorID, err := strconv.Atoi(params["minorID"])
	if err != nil {
		http.Error(w, "Beacon can't be found with invalid id", 404)
	}
	log.Println(minorID)
	beacon, err := SearchBeaconByID(minorID, 0, -1)
	log.Printf("%v", beacon)
	if err != nil {
		//Beacon not found return 404
		http.Error(w, "Beacon not found in db", 404)
	}
	for _, i := range beacon {
		arts, error := SearchArtByBeacon(i, 0, -1)
		if error != nil {
			http.Error(w, "Beacon assigned to art piece", 404)
		}
		ServeJSON(w, arts)

	}
	//	responseHeaders.Add("Accept","application/json")
	http.Error(w, "Beacon not found", 404)
	//Send a get request to api with keyword param
}

func initDB() {
	s.DBURI = DBURI
	s.dbName = dbName
	s.getSession()
	s.Session.SetSafe(&mgo.Safe{})
	s.Session.SetMode(mgo.Monotonic, true)
	cNames, errors := EnsureIndex(CollectionNames, indices...)
	for k, err := range errors {
		if err != nil {
			log.Fatalf("Indices not taking for %v;%v\n", cNames[k], err)
		}
	}
}

//EnsureIndex ensures that when we store things, we get the expected results
func EnsureIndex(collectionNames []string, indices ...mgo.Index) (s []string, e []error) {
	for j, k := range indices {
		function := func(c *mgo.Collection) error {
			return c.EnsureIndex(k)
		}
		err := withCollection(collectionNames[j], function)
		if err != nil {
			s = append(s, collectionNames[j])
			e = append(e, err)
		}
	}
	return
}

func (s *Server) getSession() *mgo.Session {
	if s.Session == nil {
		var err error
		dialInfo := &mgo.DialInfo{
			Addrs:    []string{s.DBURI},
			Direct:   true,
			FailFast: false,
		}
		s.Session, err = mgo.DialWithInfo(dialInfo)
		if err != nil {
			log.Fatalf("Can't find MongoDB. Is it started? %v\n", err)
		}
	}
	// Returns a copy of the session so we don't waste le resources? Doesn't reuse socket however
	return s.Session.Copy()
}

func withCollection(collection string, fn func(*mgo.Collection) error) error {
	session := s.getSession()
	defer session.Close()
	c := session.DB(s.dbName).C(collection)
	return fn(c)
}

//Insert datum into a specific collection
func Insert(collectionName string, values ...interface{}) error {
	function := func(c *mgo.Collection) error {
		err := c.Insert(values...)
		if err != nil {
			log.Printf("Can't insert document, %v\n", err)
		}
		return err
	}
	return withCollection(collectionName, function)
}

//SearchBeacon searches for a beacon using a passed in struct
func SearchBeacon(q interface{}, skip int, limit int) (searchResults []Beacon, err error) {
	searchResults = []Beacon{}
	query := func(c *mgo.Collection) error {
		function := c.Find(q).Skip(skip).Limit(limit).All(&searchResults)
		if limit < 0 {
			function = c.Find(q).Skip(skip).All(&searchResults)
		}
		return function
	}
	search := func() error {
		return withCollection("beacon", query)
	}
	err = search()
	return
}

//Estimote id
//

//SearchBeaconByID is a
func SearchBeaconByID(minorID int, skip int, limit int) (searchResults []Beacon, err error) {
	//	if len(beacon) == 20 {
	return SearchBeacon(bson.M{"minorid": minorID}, skip, limit)
	//	}
	//	return nil, errors.New("Not long enough to be a beacon")
}

//SearchArt is a
func SearchArt(q interface{}, skip int, limit int) (searchResults []WalterObj, err error) {
	searchResults = []WalterObj{}
	query := func(c *mgo.Collection) error {
		function := c.Find(q).Skip(skip).Limit(limit).All(&searchResults)
		if limit < 0 {
			function = c.Find(q).Skip(skip).All(&searchResults)
		}
		return function
	}
	search := func() error {
		return withCollection("art", query)
	}
	err = search()
	return
}

//SearchArtByBeacon is a specific version of searchArt
func SearchArtByBeacon(beacon Beacon, skip int, limit int) (searchResults []WalterObj, err error) {
	return SearchArt(bson.M{"beacon": beacon}, skip, limit)
}

// ServeJSON replies to the request with a JSON
// representation of resource v.
func ServeJSON(w http.ResponseWriter, v interface{}) {
	// avoid json vulnerabilities, always wrap v in an object literal
	//	doc := map[string]interface{}{"d": v}
	if data, err := json.Marshal(v); err != nil {
		log.Printf("Error marshalling json: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Header().Set("Content-Length", strconv.Itoa(len(data)))
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.Write(data)
	}
}

// ReadJSON decodes JSON data into a provided struct which must be passed in as a pointer.
//If it's not a pointer you are basically putting your data into a bottomless gorge and willing it to
//show up right next to you. Just no.
func ReadJSON(req *http.Request, v interface{}) error {
	defer req.Body.Close()
	decoder := json.NewDecoder(req.Body)
	err := decoder.Decode(v)
	return err
}

//Use this method to debug things
func logRequest(handler func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		var s = time.Now()
		handler(w, r)
		log.Printf("%s %s %6.3fms", r.Method, r.RequestURI, (time.Since(s).Seconds() * 1000))
	}
}
<commit_msg>Got Moxtra API key working<commit_after>package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/mux"
	"github.com/nu7hatch/gouuid"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

//curl -H "Content-Type: application/json" -XPUT -d '{"P"}' -f -v http://localhost:3000/beacon/{id}

const (
	dbName = "artroomServer"
	//DBURI for the mongodb
	DBURI = "127.0.0.1"

	waltersAPIPrefix    = "http://api.thewalters.org/v1/objects?apikey=ShxvahaBFNIcfWR7E78xXdssKIlXtUAJk9rDDrrmlvbOlxQKtASCzV4op5aHv2Il&title="
	waltersImagePrefix  = "http://static.thewalters.org/images/"
	waltersImagePostfix = "?width=500"
	moxtraAPI           = "https://api.moxtra.com/oauth/token?client_id="             //client-id
	moxtraCS            = "&client_secret="                                           //CLIENT-SECRET
	moxtraGrantType     = "&grant_type=http://www.moxtra.com/auth_uniqueid&uniqueid=" //UNIQUE-ID
	moxtraTS            = "&timestamp="                                               //TIMESTAMP
	moxtraSig           = "&signature="                                               //SIGNATURE
	moxtraClientID      = "u-DTusZ0ruc"
)

var (
	s = Server{}
	//CollectionNames is the shiz
	CollectionNames = []string{"beacon", "art"}
	beaconIndex     = mgo.Index{
		Key:        []string{"MinorID"},
		Unique:     true,
		DropDups:   true,
		Background: true,
		Sparse:     true,
		Name:       "beaconIndex",
		//ExpireAf
	}

	artIndex = mgo.Index{
		Key:        []string{"Beacon", "Title"},
		Unique:     true,
		DropDups:   true,
		Background: true,
		Sparse:     true,
		Name:       "ArtIndex",
		//ExpireAf
	}
	indices            = []mgo.Index{beaconIndex, artIndex}
	moxtraTimeStamp    string
	moxtraSignature    []byte
	moxtraClientSecret = []byte("vWLjlkGMZOk")
)

type (
	//Server is the name for the server deal wit it
	Server struct {
		Session     *mgo.Session // The main session we'll we be cloning
		DBURI       string       // Where the DB is on the network
		dbName      string       // Name of the MongoDB
		MoxtraToken MoxtraToken
	}

	//Beacon is the struct that structures what the data for a beacon will look like
	Beacon struct {
		ProxID  string
		MajorID int
		MinorID int
	}
	//FormResponse is what we get back from the form
	FormResponse struct {
		Date        string
		Author      string
		Title       string
		ProxID      string
		MajorID     int
		MinorID     int
		Description string
	}

	//The structure the walter api gives us more or less
	WalterObj struct {
		ObjectID       int
		Collection     string
		Title          string
		Author         string
		Medium         string
		Description    string
		Images         string
		CuratorComment string
		Beacon         Beacon
		ImageURL       string
	}
	//Walters objects
	Walters struct {
		Items []WalterObj
	}

	//MoxtraToken balls to the wall
	MoxtraToken struct {
		AccessToken string `json:"access_token"`
		TokenType   string `json:"token_type"`
		ExpiresIn   int    `json:"expires_in"`
		Scope       string
	}

	MoxtraChatRoom struct {
		Name         string
		Conversation bool
	}
)

func main() {
	log.Println("Server is warming up...")
	initDB()
	initMoxtra()
	defer s.Session.Close()
	http.Handle("/", initHandlers())
	log.Fatalln(http.ListenAndServe(":3000", nil))
}

func initMoxtra() {
	u, err := uuid.NewV4()
	if err != nil {
		log.Fatalf("Unable to Gen UUID for Moxtra.")
	}
	moxtraTimeStamp = strconv.FormatInt(time.Now().Unix()*1000, 10)
	log.Printf("%v", moxtraTimeStamp)
	combString := moxtraClientID + u.String() + moxtraTimeStamp

	mac := hmac.New(sha256.New, moxtraClientSecret)
	mac.Write([]byte(combString))
	moxtraSignature = mac.Sum(nil)
	str := base64.StdEncoding.EncodeToString(moxtraSignature)
	fmt.Println(str)
	//	re := regexp.MustCompile("/\+/g")//Global match, any backslash become a hyphen
	//	str = re.ReplaceAllLiteralString(str, "-")
	//	re = regexp.MustCompile("\\//g")//Any forward slash becomes an underscore
	//	str = re.ReplaceAllLiteralString(str, "_")
	//	re = regexp.MustCompile("\\=+$/")//Remove any trailing equals
	//	str = re.ReplaceAllLiteralString(str, "")
	str = strings.Replace(str, "\\", "-", -1)
	str = strings.Replace(str, "/", "_", -1)
	//	str = strings.Replace(str, "=", "", -1)
	str = str[:len(str)-1]
	client := &http.Client{}
	url := moxtraAPI + moxtraClientID + moxtraCS + string(moxtraClientSecret) + moxtraGrantType + u.String() + moxtraTS + moxtraTimeStamp + moxtraSig + str //string(moxtraSignature)
	log.Println(url)
	request, _ := http.NewRequest("POST", url, nil)
	//	request.Header.Add("Accept", "application/json")
	res, err := client.Do(request)
	if err != nil {
		//		http.Error(w, "Req didn't go through", 500)
		//		return
		log.Fatalf("Req didn't go through, %v", err)
	}
	moxie := MoxtraToken{}
	//err = ReadJSON(request, &moxie)
	defer res.Body.Close()
	//	var data json.RawMessage
	log.Printf("%v", res.Body)
	err = json.NewDecoder(res.Body).Decode(&moxie)
	log.Printf("\n\n%v\n", moxie)
	if err != nil {
		//		http.Error(w, "Couldn't decode data...", 500)
		//		return
		log.Fatalf("Can't resolve the datum?\n %v", err)
		//	panic(err)
	}
	log.Println("PASSED")
	//err = json.Unmarshal(data, &moxie)
	if err != nil {
		//http.Error(w, "Couldn't marshall into individual waltOBJ.", 500)
		//return
		log.Fatalf("Can't unmarshall into individ waltOBJ, %v", err)
	}
	s.MoxtraToken = moxie
	log.Printf("%v", moxie)
}

func initHandlers() *mux.Router {
	r := mux.NewRouter()
	r.HandleFunc("/beacon", handlePOSTBeacon).Methods("POST", "OPTIONS")
	r.HandleFunc("/beacon/{minorID:[0-9]+}", handleBeacon).Methods("GET")
	return r
}

func handlePOSTBeacon(w http.ResponseWriter, req *http.Request) {
	log.Println("I got here?")
	switch req.Method {
	case "OPTIONS":
		w.Header().Set("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		fallthrough
	case "POST":
		//		w.Header().Set("Access-Control-Allow-Origin", "*")
		var formResponse FormResponse
		err := ReadJSON(req, &formResponse)
		log.Printf("We have begun the ritual, %v", formResponse)
		client := &http.Client{}
		request, _ := http.NewRequest("GET", waltersAPIPrefix+strings.Replace(formResponse.Title, " ", "%20", -1), nil)
		request.Header.Add("Accept", "application/json")
		res, err := client.Do(request)
		if err != nil {
			http.Error(w, "Req didn't go through", 500)
			return
			//		log.Fatalf("Req didn't go through, %v", err)
		}
		defer res.Body.Close()
		var data json.RawMessage
		err = json.NewDecoder(res.Body).Decode(&data)
		if err != nil {
			http.Error(w, "Couldn't decode data...", 500)
			return
			//		log.Fatalf("Can't resolve the datum?%v", err)
		}
		wal := Walters{}
		err = json.Unmarshal(data, &wal)
		if err != nil {
			http.Error(w, "Couldn't marshall into individual waltOBJ.", 500)
			return
			//		log.Fatalf("Can't unmarshall into individ waltOBJ, %v", err)
		}
		if len(wal.Items) < 1 {
			http.Error(w, "Couldn't marshall correctly", 500)
			return
		}
		waltersObject := wal.Items[0]
		waltersObject.Beacon = Beacon{formResponse.ProxID, formResponse.MajorID, formResponse.MinorID}
		waltersObject.ImageURL = waltersImagePrefix + string(waltersObject.Images[0]) + waltersImagePostfix
		waltersObject.CuratorComment = formResponse.Description
		err = Insert("beacon", waltersObject.Beacon)
		if err != nil {
			//Return a 500
			http.Error(w, "Couldn't insert corresponding beacon.", 500)
			return
		}

		err = Insert("art", waltersObject)
		if err != nil {
			//Return a 500
			http.Error(w, "Couldn't insert corresponding art piece.", 500)
			return
		}
		//https//api.moxtra.com/me/binders
		http.Error(w, http.StatusText(http.StatusCreated), http.StatusCreated)
		break
	}
}

func handleBeacon(w http.ResponseWriter, req *http.Request) {
	params := mux.Vars(req)
	minorID, err := strconv.Atoi(params["minorID"])
	if err != nil {
		http.Error(w, "Beacon can't be found with invalid id", 404)
	}
	log.Println(minorID)
	beacon, err := SearchBeaconByID(minorID, 0, -1)
	log.Printf("%v", beacon)
	if err != nil {
		//Beacon not found return 404
		http.Error(w, "Beacon not found in db", 404)
	}
	for _, i := range beacon {
		arts, error := SearchArtByBeacon(i, 0, -1)
		if error != nil {
			http.Error(w, "Beacon assigned to art piece", 404)
		}
		ServeJSON(w, arts)

	}
	//	responseHeaders.Add("Accept","application/json")
	http.Error(w, "Beacon not found", 404)
	//Send a get request to api with keyword param
}

func initDB() {
	s.DBURI = DBURI
	s.dbName = dbName
	s.getSession()
	s.Session.SetSafe(&mgo.Safe{})
	s.Session.SetMode(mgo.Monotonic, true)
	cNames, errors := EnsureIndex(CollectionNames, indices...)
	for k, err := range errors {
		if err != nil {
			log.Fatalf("Indices not taking for %v;%v\n", cNames[k], err)
		}
	}
}

//EnsureIndex ensures that when we store things, we get the expected results
func EnsureIndex(collectionNames []string, indices ...mgo.Index) (s []string, e []error) {
	for j, k := range indices {
		function := func(c *mgo.Collection) error {
			return c.EnsureIndex(k)
		}
		err := withCollection(collectionNames[j], function)
		if err != nil {
			s = append(s, collectionNames[j])
			e = append(e, err)
		}
	}
	return
}

func (s *Server) getSession() *mgo.Session {
	if s.Session == nil {
		var err error
		dialInfo := &mgo.DialInfo{
			Addrs:    []string{s.DBURI},
			Direct:   true,
			FailFast: false,
		}
		s.Session, err = mgo.DialWithInfo(dialInfo)
		if err != nil {
			log.Fatalf("Can't find MongoDB. Is it started? %v\n", err)
		}
	}
	// Returns a copy of the session so we don't waste le resources? Doesn't reuse socket however
	return s.Session.Copy()
}

func withCollection(collection string, fn func(*mgo.Collection) error) error {
	session := s.getSession()
	defer session.Close()
	c := session.DB(s.dbName).C(collection)
	return fn(c)
}

//Insert datum into a specific collection
func Insert(collectionName string, values ...interface{}) error {
	function := func(c *mgo.Collection) error {
		err := c.Insert(values...)
		if err != nil {
			log.Printf("Can't insert document, %v\n", err)
		}
		return err
	}
	return withCollection(collectionName, function)
}

//SearchBeacon searches for a beacon using a passed in struct
func SearchBeacon(q interface{}, skip int, limit int) (searchResults []Beacon, err error) {
	searchResults = []Beacon{}
	query := func(c *mgo.Collection) error {
		function := c.Find(q).Skip(skip).Limit(limit).All(&searchResults)
		if limit < 0 {
			function = c.Find(q).Skip(skip).All(&searchResults)
		}
		return function
	}
	search := func() error {
		return withCollection("beacon", query)
	}
	err = search()
	return
}

//Estimote id
//

//SearchBeaconByID is a
func SearchBeaconByID(minorID int, skip int, limit int) (searchResults []Beacon, err error) {
	//	if len(beacon) == 20 {
	return SearchBeacon(bson.M{"minorid": minorID}, skip, limit)
	//	}
	//	return nil, errors.New("Not long enough to be a beacon")
}

//SearchArt is a
func SearchArt(q interface{}, skip int, limit int) (searchResults []WalterObj, err error) {
	searchResults = []WalterObj{}
	query := func(c *mgo.Collection) error {
		function := c.Find(q).Skip(skip).Limit(limit).All(&searchResults)
		if limit < 0 {
			function = c.Find(q).Skip(skip).All(&searchResults)
		}
		return function
	}
	search := func() error {
		return withCollection("art", query)
	}
	err = search()
	return
}

//SearchArtByBeacon is a specific version of searchArt
func SearchArtByBeacon(beacon Beacon, skip int, limit int) (searchResults []WalterObj, err error) {
	return SearchArt(bson.M{"beacon": beacon}, skip, limit)
}

// ServeJSON replies to the request with a JSON
// representation of resource v.
func ServeJSON(w http.ResponseWriter, v interface{}) {
	// avoid json vulnerabilities, always wrap v in an object literal
	//	doc := map[string]interface{}{"d": v}
	if data, err := json.Marshal(v); err != nil {
		log.Printf("Error marshalling json: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Header().Set("Content-Length", strconv.Itoa(len(data)))
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.Write(data)
	}
}

// ReadJSON decodes JSON data into a provided struct which must be passed in as a pointer.
//If it's not a pointer you are basically putting your data into a bottomless gorge and willing it to
//show up right next to you. Just no.
func ReadJSON(req *http.Request, v interface{}) error {
	defer req.Body.Close()
	decoder := json.NewDecoder(req.Body)
	err := decoder.Decode(v)
	return err
}

//Use this method to debug things
func logRequest(handler func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		var s = time.Now()
		handler(w, r)
		log.Printf("%s %s %6.3fms", r.Method, r.RequestURI, (time.Since(s).Seconds() * 1000))
	}
}
<|endoftext|>


<commit_before>package gontpd

import (
	"context"
	"fmt"
	"log"
	"net"
	"syscall"
	"time"

	"golang.org/x/sys/unix"
)

const limit = 4

func (d *NTPd) listen() {

	for j := 0; j < d.cfg.ConnNum; j++ {
		conn, err := d.makeConn()
		if err != nil {
			log.Fatal(err)
		}
		for i := 0; i < d.cfg.WorkerNum; i++ {
			id := fmt.Sprintf("%d:%d", j, i)
			w := worker{
				id, newLRU(d.cfg.CacheSize),
				conn, d.stat, d,
			}
			go w.Work()
		}
	}
}

type worker struct {
	id   string
	lru  *lru
	conn *net.UDPConn
	stat *statistic
	d    *NTPd
}

func (d *NTPd) makeConn() (conn *net.UDPConn, err error) {

	var operr error

	cfgFn := func(network, address string, conn syscall.RawConn) (err error) {

		fn := func(fd uintptr) {
			operr = syscall.SetsockoptInt(int(fd),
				syscall.SOL_SOCKET,
				unix.SO_REUSEPORT, 1)
			if operr != nil {
				return
			}
		}

		if err = conn.Control(fn); err != nil {
			return err
		}

		err = operr
		return
	}
	lc := net.ListenConfig{Control: cfgFn}
	lp, err := lc.ListenPacket(context.Background(), "udp", d.cfg.Listen)
	if err != nil {
		return
	}
	conn = lp.(*net.UDPConn)
	return
}

func (w *worker) Work() {
	var (
		receiveTime time.Time
		remoteAddr  *net.UDPAddr

		err      error
		addrS    string
		lastUnix int64
		n        int
		ok       bool
	)
	p := make([]byte, 48)
	oob := make([]byte, 1)

	defer func() {
		if r := recover(); r != nil {
			log.Printf("Worker: %s fatal, reason:%s, read:%d", w.id, r, n)
		} else {
			log.Printf("Worker: %s exited, reason:%s, read:%d", w.id, err, n)
		}
	}()

	log.Printf("worker %s started", w.id)

	for {
		n, _, _, remoteAddr, err = w.conn.ReadMsgUDP(p, oob)
		if err != nil {
			return
		}

		receiveTime = time.Now()
		if n < 48 {
			if debug {
				log.Printf("worker: %s get small packet %d",
					remoteAddr.String(), n)
			}
			if w.stat != nil {
				w.stat.fastDropCounter.WithLabelValues("small_packet").Inc()
			}
			continue
		}

		if w.d.cfg.LanDrop && isLan(remoteAddr.IP) {
			if debug {
				log.Printf("worker: %s get lan dest",
					remoteAddr.String())
			}
			if w.stat != nil {
				w.stat.fastDropCounter.WithLabelValues("lan_dest").Inc()
			}
			continue
		}

		addrS = remoteAddr.IP.String()
		lastUnix, ok = w.lru.Get(addrS)
		if ok && receiveTime.Unix()-lastUnix < limit {
			if debug {
				log.Println(receiveTime.Unix()-lastUnix, ok)
			}
			w.sendError(remoteAddr, rateKoD)
			if w.stat != nil {
				w.stat.fastDropCounter.WithLabelValues("rate").Inc()
			}
			continue
		}

		// BCE
		_ = p[47]

		w.lru.Add(addrS, receiveTime.Unix())

		// GetMode

		switch p[LiVnModePos] &^ 0xf8 {
		case ModeSymmetricActive:
			w.sendError(remoteAddr, acstKoD)
		case ModeReserved:
			fallthrough
		case ModeClient:
			copy(p[0:OriginTimeStamp], w.d.template)
			copy(p[OriginTimeStamp:OriginTimeStamp+8],
				p[TransmitTimeStamp:TransmitTimeStamp+8])
			SetUint64(p, ReceiveTimeStamp, toNtpTime(receiveTime))

			SetUint64(p, TransmitTimeStamp, toNtpTime(time.Now()))
			_, err = w.conn.WriteToUDP(p, remoteAddr)
			if err != nil && debug {
				log.Printf("worker: %s write failed. %s", remoteAddr.String(), err)
			}
			if w.stat != nil {
				w.stat.fastCounter.Inc()
				w.logIP(remoteAddr)
			}
		default:
			if debug {
				log.Printf("%s not support client request mode:%x",
					remoteAddr.String(), p[LiVnModePos]&^0xf8)
			}
			if w.stat != nil {
				w.stat.fastDropCounter.WithLabelValues("unknown_mode").Inc()
			}
		}
	}
}

func (w *worker) sendError(raddr *net.UDPAddr, err uint32) {
	// return
	errBuf := make([]byte, 48)
	copy(errBuf, w.d.template)
	SetUint8(errBuf, StratumPos, 0)
	SetUint32(errBuf, ReferIDPos, err)
	w.conn.WriteToUDP(errBuf, raddr)
}

func (w *worker) logIP(raddr *net.UDPAddr) {
	if w.stat.geoDB == nil {
		return
	}
	country, err := w.stat.geoDB.LookupIP(raddr.IP)
	if err != nil {
		return
	}
	if country.Country == nil {
		return
	}
	w.stat.reqCounter.WithLabelValues(country.Country.Code).Inc()
}
<commit_msg>fix rate<commit_after>package gontpd

import (
	"context"
	"fmt"
	"log"
	"net"
	"syscall"
	"time"

	"golang.org/x/sys/unix"
)

// minimum headway time is 2 seconds, https://www.eecis.udel.edu/~mills/ntp/html/rate.html
const limit = 2

func (d *NTPd) listen() {

	for j := 0; j < d.cfg.ConnNum; j++ {
		conn, err := d.makeConn()
		if err != nil {
			log.Fatal(err)
		}
		for i := 0; i < d.cfg.WorkerNum; i++ {
			id := fmt.Sprintf("%d:%d", j, i)
			w := worker{
				id, newLRU(d.cfg.CacheSize),
				conn, d.stat, d,
			}
			go w.Work()
		}
	}
}

type worker struct {
	id   string
	lru  *lru
	conn *net.UDPConn
	stat *statistic
	d    *NTPd
}

func (d *NTPd) makeConn() (conn *net.UDPConn, err error) {

	var operr error

	cfgFn := func(network, address string, conn syscall.RawConn) (err error) {

		fn := func(fd uintptr) {
			operr = syscall.SetsockoptInt(int(fd),
				syscall.SOL_SOCKET,
				unix.SO_REUSEPORT, 1)
			if operr != nil {
				return
			}
		}

		if err = conn.Control(fn); err != nil {
			return err
		}

		err = operr
		return
	}
	lc := net.ListenConfig{Control: cfgFn}
	lp, err := lc.ListenPacket(context.Background(), "udp", d.cfg.Listen)
	if err != nil {
		return
	}
	conn = lp.(*net.UDPConn)
	return
}

func (w *worker) Work() {
	var (
		receiveTime time.Time
		remoteAddr  *net.UDPAddr

		err      error
		addrS    string
		lastUnix int64
		n        int
		ok       bool
	)
	p := make([]byte, 48)
	oob := make([]byte, 1)

	defer func() {
		if r := recover(); r != nil {
			log.Printf("Worker: %s fatal, reason:%s, read:%d", w.id, r, n)
		} else {
			log.Printf("Worker: %s exited, reason:%s, read:%d", w.id, err, n)
		}
	}()

	log.Printf("worker %s started", w.id)

	for {
		n, _, _, remoteAddr, err = w.conn.ReadMsgUDP(p, oob)
		if err != nil {
			return
		}

		receiveTime = time.Now()
		if n < 48 {
			if debug {
				log.Printf("worker: %s get small packet %d",
					remoteAddr.String(), n)
			}
			if w.stat != nil {
				w.stat.fastDropCounter.WithLabelValues("small_packet").Inc()
			}
			continue
		}

		if w.d.cfg.LanDrop && isLan(remoteAddr.IP) {
			if debug {
				log.Printf("worker: %s get lan dest",
					remoteAddr.String())
			}
			if w.stat != nil {
				w.stat.fastDropCounter.WithLabelValues("lan_dest").Inc()
			}
			continue
		}

		// BCE
		_ = p[47]

		addrS = remoteAddr.IP.String()
		lastUnix, ok = w.lru.Get(addrS)
		if ok && receiveTime.Unix()-lastUnix < limit {
			if debug {
				log.Println(receiveTime.Unix()-lastUnix, ok)
			}
			w.sendError(p, remoteAddr, rateKoD)
			if w.stat != nil {
				w.stat.fastDropCounter.WithLabelValues("rate").Inc()
			}
			continue
		}

		w.lru.Add(addrS, receiveTime.Unix())

		// GetMode

		switch p[LiVnModePos] &^ 0xf8 {
		case ModeSymmetricActive:
			w.sendError(p, remoteAddr, acstKoD)
		case ModeReserved:
			fallthrough
		case ModeClient:
			copy(p[0:OriginTimeStamp], w.d.template)
			copy(p[OriginTimeStamp:OriginTimeStamp+8],
				p[TransmitTimeStamp:TransmitTimeStamp+8])
			SetUint64(p, ReceiveTimeStamp, toNtpTime(receiveTime))

			SetUint64(p, TransmitTimeStamp, toNtpTime(time.Now()))
			_, err = w.conn.WriteToUDP(p, remoteAddr)
			if err != nil && debug {
				log.Printf("worker: %s write failed. %s", remoteAddr.String(), err)
			}
			if w.stat != nil {
				w.stat.fastCounter.Inc()
				w.logIP(remoteAddr)
			}
		default:
			if debug {
				log.Printf("%s not support client request mode:%x",
					remoteAddr.String(), p[LiVnModePos]&^0xf8)
			}
			if w.stat != nil {
				w.stat.fastDropCounter.WithLabelValues("unknown_mode").Inc()
			}
		}
	}
}

func (w *worker) sendError(p []byte, raddr *net.UDPAddr, err uint32) {
	// avoid spoof
	copy(p[0:OriginTimeStamp], w.d.template)
	copy(p[OriginTimeStamp:OriginTimeStamp+8],
		p[TransmitTimeStamp:TransmitTimeStamp+8])
	SetUint8(p, StratumPos, 0)
	SetUint32(p, ReferIDPos, err)
	w.conn.WriteToUDP(p, raddr)
}

func (w *worker) logIP(raddr *net.UDPAddr) {
	if w.stat.geoDB == nil {
		return
	}
	country, err := w.stat.geoDB.LookupIP(raddr.IP)
	if err != nil {
		return
	}
	if country.Country == nil {
		return
	}
	w.stat.reqCounter.WithLabelValues(country.Country.Code).Inc()
}
<|endoftext|>


<commit_before>package ipfix

import (
	"time"
	"fmt"
	"net/http"
	"os"
	"runtime/debug"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"github.com/go-chi/chi/middleware"
)

type middlewareHandler = func(next http.Handler) http.Handler

type recoverMiddleware struct {
	debug bool
	logger *zap.Logger
}

func newRecoverMiddleware(debug bool, logger *zap.Logger) *recoverMiddleware {
	return &recoverMiddleware{
		debug: debug,
		logger: logger,
	}
}

func (m *recoverMiddleware) Handle(it interface{}) {
	if m.debug {
		fmt.Fprintf(os.Stderr, "Panic: %+v\n", it)
		debug.PrintStack()
	} else {
		m.logger.Error("Error handled")
	}
}

func (m *recoverMiddleware) Handler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if rvr := recover(); rvr != nil {
				m.Handle(rvr)

				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			}
		}()

		next.ServeHTTP(w, r)
	})
}

func newLoggerMiddleware(logger *zap.Logger) func(next http.Handler) http.Handler {
	return middleware.RequestLogger(&StructuredLogger{logger})
}

type StructuredLogger struct {
	Logger *zap.Logger
}

func (l *StructuredLogger) NewLogEntry(r *http.Request) middleware.LogEntry {
	entry := &StructuredLoggerEntry{Logger: l.Logger}

	fields := []zapcore.Field{zap.String("ts", time.Now().UTC().Format(time.RFC1123))}

	if reqID := middleware.GetReqID(r.Context()); reqID != "" {
		fields = append(fields, zap.String("req.id", reqID))
	}

	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}

	fields = append(fields, []zapcore.Field{
		zap.String("http.scheme", scheme),
		zap.String("http.proto", r.Proto),
		zap.String("http.method", r.Method),
		zap.String("remote_addr", r.RemoteAddr),
		zap.String("user_agent", r.UserAgent()),
		zap.String("uri", fmt.Sprintf("%s://%s%s", scheme, r.Host, r.RequestURI)),
	}...)

	entry.Logger = l.Logger.With(fields...)

	entry.Logger.Info("request started")

	return entry
}

type StructuredLoggerEntry struct {
	Logger *zap.Logger
}

func (l *StructuredLoggerEntry) Write(status, bytes int, elapsed time.Duration) {
	l.Logger = l.Logger.With(
		zap.Int("res.status", status),
		zap.Int("res.bytes_length", bytes),
		zap.Float64("res.elapsed_ms", float64(elapsed.Nanoseconds()) / 1000000.0))

	l.Logger.Info("request complete")
}

func (l *StructuredLoggerEntry) Panic(v interface{}, stack []byte) {
	l.Logger = l.Logger.With(
		zap.String("stack", string(stack)),
		zap.String("panic", fmt.Sprintf("%+v", v)),
	)
}
<commit_msg>chore: middleware should stay private<commit_after>package ipfix

import (
	"time"
	"fmt"
	"net/http"
	"os"
	"runtime/debug"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"github.com/go-chi/chi/middleware"
)

type middlewareHandler = func(next http.Handler) http.Handler

type recoverMiddleware struct {
	debug bool
	logger *zap.Logger
}

func newRecoverMiddleware(debug bool, logger *zap.Logger) *recoverMiddleware {
	return &recoverMiddleware{
		debug: debug,
		logger: logger,
	}
}

func (m *recoverMiddleware) Handle(it interface{}) {
	if m.debug {
		fmt.Fprintf(os.Stderr, "Panic: %+v\n", it)
		debug.PrintStack()
	} else {
		m.logger.Error("Error handled")
	}
}

func (m *recoverMiddleware) Handler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if rvr := recover(); rvr != nil {
				m.Handle(rvr)

				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			}
		}()

		next.ServeHTTP(w, r)
	})
}

func newLoggerMiddleware(logger *zap.Logger) func(next http.Handler) http.Handler {
	return middleware.RequestLogger(&structuredLogger{logger})
}

type structuredLogger struct {
	Logger *zap.Logger
}

func (l *structuredLogger) NewLogEntry(r *http.Request) middleware.LogEntry {
	entry := &structuredLoggerEntry{Logger: l.Logger}

	fields := []zapcore.Field{zap.String("ts", time.Now().UTC().Format(time.RFC1123))}

	if reqID := middleware.GetReqID(r.Context()); reqID != "" {
		fields = append(fields, zap.String("req.id", reqID))
	}

	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}

	fields = append(fields, []zapcore.Field{
		zap.String("http.scheme", scheme),
		zap.String("http.proto", r.Proto),
		zap.String("http.method", r.Method),
		zap.String("remote_addr", r.RemoteAddr),
		zap.String("user_agent", r.UserAgent()),
		zap.String("uri", fmt.Sprintf("%s://%s%s", scheme, r.Host, r.RequestURI)),
	}...)

	entry.Logger = l.Logger.With(fields...)

	entry.Logger.Info("request started")

	return entry
}

type structuredLoggerEntry struct {
	Logger *zap.Logger
}

func (l *structuredLoggerEntry) Write(status, bytes int, elapsed time.Duration) {
	l.Logger = l.Logger.With(
		zap.Int("res.status", status),
		zap.Int("res.bytes_length", bytes),
		zap.Float64("res.elapsed_ms", float64(elapsed.Nanoseconds()) / 1000000.0))

	l.Logger.Info("request complete")
}

func (l *structuredLoggerEntry) Panic(v interface{}, stack []byte) {
	l.Logger = l.Logger.With(
		zap.String("stack", string(stack)),
		zap.String("panic", fmt.Sprintf("%+v", v)),
	)
}
<|endoftext|>