<commit_before>// An ESMTP server library.
package smtp

import (
	"crypto/tls"
	"errors"
	"io"
	"net"

	"github.com/emersion/go-sasl"
)

// A function that creates SASL servers.
type SaslServerFactory func(conn *Conn) sasl.Server

// A SMTP server.
type Server struct {
	// TCP address to listen on.
	Addr string
	// The server TLS configuration.
	TLSConfig *tls.Config

	Domain            string
	MaxRecipients     int
	MaxIdleSeconds    int
	MaxMessageBytes   int
	AllowInsecureAuth bool
	Debug             io.Writer

	// The server backend.
	Backend Backend

	listener net.Listener
	caps     []string
	auths    map[string]SaslServerFactory
}

// New creates a new SMTP server.
func New(bkd Backend) *Server {
	return &Server{
		Backend: bkd,
		caps:    []string{"PIPELINING", "8BITMIME"},
		auths: map[string]SaslServerFactory{
			"PLAIN": func(conn *Conn) sasl.Server {
				return sasl.NewPlainServer(func(identity, username, password string) error {
					if identity != "" && identity != username {
						return errors.New("Identities not supported")
					}

					user, err := bkd.Login(username, password)
					if err != nil {
						return err
					}

					conn.User = user
					return nil
				})
			},
		},
	}
}

// Serve accepts incoming connections on the Listener l.
func (s *Server) Serve(l net.Listener) error {
	s.listener = l
	defer s.Close()

	for {
		c, err := l.Accept()
		if err != nil {
			return err
		}

		go s.handleConn(newConn(c, s))
	}
}

func (s *Server) handleConn(c *Conn) error {
	defer c.Close()
	c.greet()

	for {
		line, err := c.readLine()
		if err == nil {
			cmd, arg, err := parseCmd(line)
			if err != nil {
				c.nbrErrors++
				c.Write("501", "Bad command")
				continue
			}

			c.handle(cmd, arg)
		} else {
			if err == io.EOF {
				return nil
			}

			if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
				c.Write("221", "Idle timeout, bye bye")
				return nil
			}

			c.Write("221", "Connection error, sorry")
			return err
		}
	}
}

// ListenAndServe listens on the TCP network address s.Addr and then calls Serve
// to handle requests on incoming connections.
//
// If s.Addr is blank, ":smtp" is used.
func (s *Server) ListenAndServe() error {
	addr := s.Addr
	if addr == "" {
		addr = ":smtp"
	}

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

	return s.Serve(l)
}

// ListenAndServeTLS listens on the TCP network address s.Addr and then calls
// Serve to handle requests on incoming TLS connections.
//
// If s.Addr is blank, ":smtps" is used.
func (s *Server) ListenAndServeTLS() error {
	addr := s.Addr
	if addr == "" {
		addr = ":smtps"
	}

	l, err := tls.Listen("tcp", addr, s.TLSConfig)
	if err != nil {
		return err
	}

	return s.Serve(l)
}

// Close stops the server.
func (s *Server) Close() {
	// TODO: say bye to all clients

	s.listener.Close()
}

// EnableAuth enables an authentication mechanism on this server.
//
// This function should not be called directly, it must only be used by
// libraries implementing extensions of the IMAP protocol.
func (s *Server) EnableAuth(name string, f SaslServerFactory) {
	s.auths[name] = f
}
<commit_msg>Use constant from sasl instead of hardcoding auth mechanism name<commit_after>// An ESMTP server library.
package smtp

import (
	"crypto/tls"
	"errors"
	"io"
	"net"

	"github.com/emersion/go-sasl"
)

// A function that creates SASL servers.
type SaslServerFactory func(conn *Conn) sasl.Server

// A SMTP server.
type Server struct {
	// TCP address to listen on.
	Addr string
	// The server TLS configuration.
	TLSConfig *tls.Config

	Domain            string
	MaxRecipients     int
	MaxIdleSeconds    int
	MaxMessageBytes   int
	AllowInsecureAuth bool
	Debug             io.Writer

	// The server backend.
	Backend Backend

	listener net.Listener
	caps     []string
	auths    map[string]SaslServerFactory
}

// New creates a new SMTP server.
func New(bkd Backend) *Server {
	return &Server{
		Backend: bkd,
		caps:    []string{"PIPELINING", "8BITMIME"},
		auths: map[string]SaslServerFactory{
			sasl.Plain: func(conn *Conn) sasl.Server {
				return sasl.NewPlainServer(func(identity, username, password string) error {
					if identity != "" && identity != username {
						return errors.New("Identities not supported")
					}

					user, err := bkd.Login(username, password)
					if err != nil {
						return err
					}

					conn.User = user
					return nil
				})
			},
		},
	}
}

// Serve accepts incoming connections on the Listener l.
func (s *Server) Serve(l net.Listener) error {
	s.listener = l
	defer s.Close()

	for {
		c, err := l.Accept()
		if err != nil {
			return err
		}

		go s.handleConn(newConn(c, s))
	}
}

func (s *Server) handleConn(c *Conn) error {
	defer c.Close()
	c.greet()

	for {
		line, err := c.readLine()
		if err == nil {
			cmd, arg, err := parseCmd(line)
			if err != nil {
				c.nbrErrors++
				c.Write("501", "Bad command")
				continue
			}

			c.handle(cmd, arg)
		} else {
			if err == io.EOF {
				return nil
			}

			if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
				c.Write("221", "Idle timeout, bye bye")
				return nil
			}

			c.Write("221", "Connection error, sorry")
			return err
		}
	}
}

// ListenAndServe listens on the TCP network address s.Addr and then calls Serve
// to handle requests on incoming connections.
//
// If s.Addr is blank, ":smtp" is used.
func (s *Server) ListenAndServe() error {
	addr := s.Addr
	if addr == "" {
		addr = ":smtp"
	}

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

	return s.Serve(l)
}

// ListenAndServeTLS listens on the TCP network address s.Addr and then calls
// Serve to handle requests on incoming TLS connections.
//
// If s.Addr is blank, ":smtps" is used.
func (s *Server) ListenAndServeTLS() error {
	addr := s.Addr
	if addr == "" {
		addr = ":smtps"
	}

	l, err := tls.Listen("tcp", addr, s.TLSConfig)
	if err != nil {
		return err
	}

	return s.Serve(l)
}

// Close stops the server.
func (s *Server) Close() {
	// TODO: say bye to all clients

	s.listener.Close()
}

// EnableAuth enables an authentication mechanism on this server.
//
// This function should not be called directly, it must only be used by
// libraries implementing extensions of the IMAP protocol.
func (s *Server) EnableAuth(name string, f SaslServerFactory) {
	s.auths[name] = f
}
<|endoftext|>


<commit_before>package mithril

import (
	"encoding/base64"
	"fmt"
	"mithril/log"
	"mithril/message"
	"mithril/store"
	"net/http"

	_ "net/http/pprof" // hey, why not
)

const faviconBase64 = `
AAABAAEAEBAAAAEAIABoBAAAFgAAACgAAAAQAAAAIAAAAAEAIAAAAAAAAAQAABILAAASCw
AAAAAAAAAAAAD//////////zMna/8zJ2v/Mydr/zMna/8zJ2v/////////////////////
/////////////////////////////////zMna/8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna/
///////////////////////////////////////////zMna/8zJ2v/Mydr////////////
/////zMna/8zJ2v/Mydr//////////////////////////////////////8zJ2v/Mydr//
//////////////////////////Mydr/zMna///////////////////////////////////
////////////////////Mydr/zMna////////////zMna/8zJ2v///////////8zJ2v/My
dr//////////////////////////////////////8zJ2v/Mydr//////8zJ2v/Mydr////
//8zJ2v/Mydr/////////////////////////////////////////////////zMna/8zJ2
v/Mydr/zMna/8zJ2v/Mydr////////////////////////////////////////////////
//////8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna/////////////////////////////////
////////////////8zJ2v/Mydr//////8zJ2v/Mydr//////8zJ2v/Mydr////////////
////////////////////////////////Mydr////////////Mydr/zMna////////////z
Mna////////////////////////////////////////////zMna////////////zMna/8z
J2v///////////8zJ2v/////////////////////////////////////////////////My
dr/zMna/8zJ2v/Mydr/zMna/8zJ2v/////////////////////////////////////////
////////////////////////Mydr/zMna/////////////////////////////////////
//Mydr/zMna/8zJ2v//////////////////////zMna/8zJ2v//////zMna/8zJ2v/Mydr
/zMna/8zJ2v///////////8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna/
8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna////////////zMna/8zJ2v/Mydr/zMna/8zJ2v/
//////////////////////////////////////////8zJ2v/AAAAAAAAAAAAAAAAAAAAAA
AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==
`

var faviconBytes []byte

type Server struct {
	amqp    *AMQPPublisher
	storage *store.Storage
	address string
}

func init() {
	faviconBytes, _ = base64.StdEncoding.DecodeString(faviconBase64)
}

func NewServer(configuration *Configuration) (*Server, error) {
	var (
		storer *store.Storage
		amqp   *AMQPPublisher
		err    error
	)

	if storer, err = store.Open(configuration.Storage, configuration.StorageUri); err != nil {
		return nil, err
	}

	if amqp, err = NewAMQPPublisher(configuration.AmqpUri); err != nil {
		return nil, err
	}

	return &Server{
		storage: storer,
		amqp:    amqp,
		address: configuration.ServerAddress,
	}, nil
}

func (me *Server) Serve() {
	http.Handle("/", me)
	log.Println("Serving on", me.address)
	log.Fatal(http.ListenAndServe(me.address, nil))
}

func (me *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	if r.Method == "POST" || r.Method == "PUT" {
		me.processMessage(w, r)
		return
	}

	if r.Method == "GET" && r.URL.Path == "/favicon.ico" {
		me.respondFavicon(w)
		return
	}

	me.respondErr(
		fmt.Errorf(`Only "POST" and "PUT" are accepted, not %q`, r.Method),
		http.StatusMethodNotAllowed,
		w)
}

func (me *Server) processMessage(w http.ResponseWriter, r *http.Request) {

	var (
		msg *message.Message
		err error
	)
	if msg, err = message.NewMessage(r); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	log.Println("Processing message: ", msg.MessageId)
	if err = me.storage.Store(msg); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	if err = me.amqp.Publish(msg); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	w.WriteHeader(http.StatusNoContent)
	w.Write([]byte(""))
}

func (me *Server) respondErr(err error, status int, w http.ResponseWriter) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(status)
	fmt.Fprintf(w, "WOMP WOMP: %v\n", err)
}

func (me *Server) respondFavicon(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "image/vnd.microsoft.icon")
	w.WriteHeader(http.StatusOK)
	w.Write(faviconBytes)
}
<commit_msg>Adding logging around request type, url, and time to handle<commit_after>package mithril

import (
	"encoding/base64"
	"fmt"
	"mithril/log"
	"mithril/message"
	"mithril/store"
	"net/http"
	"time"

	_ "net/http/pprof" // hey, why not
)

const faviconBase64 = `
AAABAAEAEBAAAAEAIABoBAAAFgAAACgAAAAQAAAAIAAAAAEAIAAAAAAAAAQAABILAAASCw
AAAAAAAAAAAAD//////////zMna/8zJ2v/Mydr/zMna/8zJ2v/////////////////////
/////////////////////////////////zMna/8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna/
///////////////////////////////////////////zMna/8zJ2v/Mydr////////////
/////zMna/8zJ2v/Mydr//////////////////////////////////////8zJ2v/Mydr//
//////////////////////////Mydr/zMna///////////////////////////////////
////////////////////Mydr/zMna////////////zMna/8zJ2v///////////8zJ2v/My
dr//////////////////////////////////////8zJ2v/Mydr//////8zJ2v/Mydr////
//8zJ2v/Mydr/////////////////////////////////////////////////zMna/8zJ2
v/Mydr/zMna/8zJ2v/Mydr////////////////////////////////////////////////
//////8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna/////////////////////////////////
////////////////8zJ2v/Mydr//////8zJ2v/Mydr//////8zJ2v/Mydr////////////
////////////////////////////////Mydr////////////Mydr/zMna////////////z
Mna////////////////////////////////////////////zMna////////////zMna/8z
J2v///////////8zJ2v/////////////////////////////////////////////////My
dr/zMna/8zJ2v/Mydr/zMna/8zJ2v/////////////////////////////////////////
////////////////////////Mydr/zMna/////////////////////////////////////
//Mydr/zMna/8zJ2v//////////////////////zMna/8zJ2v//////zMna/8zJ2v/Mydr
/zMna/8zJ2v///////////8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna/
8zJ2v/Mydr/zMna/8zJ2v/Mydr/zMna////////////zMna/8zJ2v/Mydr/zMna/8zJ2v/
//////////////////////////////////////////8zJ2v/AAAAAAAAAAAAAAAAAAAAAA
AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==
`

var faviconBytes []byte

type Server struct {
	amqp    *AMQPPublisher
	storage *store.Storage
	address string
}

func init() {
	faviconBytes, _ = base64.StdEncoding.DecodeString(faviconBase64)
}

func NewServer(configuration *Configuration) (*Server, error) {
	var (
		storer *store.Storage
		amqp   *AMQPPublisher
		err    error
	)

	if storer, err = store.Open(configuration.Storage, configuration.StorageUri); err != nil {
		return nil, err
	}

	if amqp, err = NewAMQPPublisher(configuration.AmqpUri); err != nil {
		return nil, err
	}

	return &Server{
		storage: storer,
		amqp:    amqp,
		address: configuration.ServerAddress,
	}, nil
}

func (me *Server) Serve() {
	http.Handle("/", me)
	log.Println("Serving on", me.address)
	log.Fatal(http.ListenAndServe(me.address, nil))
}

func (me *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()

	log.Printf("%s %s", r.Method, r.URL.Path)

	if r.Method == "POST" || r.Method == "PUT" {
		me.processMessage(w, r)
	} else if r.Method == "GET" && r.URL.Path == "/favicon.ico" {
		me.respondFavicon(w)
	} else {
		me.respondErr(
			fmt.Errorf(`Only "POST" and "PUT" are accepted, not %q`, r.Method),
			http.StatusMethodNotAllowed,
			w)
	}

	log.Printf("Request handled in %s.\n", time.Now().Sub(startTime))
	return
}

func (me *Server) processMessage(w http.ResponseWriter, r *http.Request) {
	var (
		msg *message.Message
		err error
	)
	if msg, err = message.NewMessage(r); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	log.Println("Processing message: ", msg.MessageId)
	if err = me.storage.Store(msg); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	if err = me.amqp.Publish(msg); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	w.WriteHeader(http.StatusNoContent)
	w.Write([]byte(""))
}

func (me *Server) respondErr(err error, status int, w http.ResponseWriter) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(status)
	fmt.Fprintf(w, "WOMP WOMP: %v\n", err)
}

func (me *Server) respondFavicon(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "image/vnd.microsoft.icon")
	w.WriteHeader(http.StatusOK)
	w.Write(faviconBytes)
}
<|endoftext|>


<commit_before>package ddns_server

import (
	log "github.com/Sirupsen/logrus"
	"ddns-server/database"
)

func main() {
	cfg, err := GetConfig()
	if err != nil {
		log.Fatal("Can't read config values")
	}

	db, err := database.New(cfg.DbHost, cfg.Port, cfg.DbName, cfg.DbUser, cfg.DbPassword)
	if err != nil {
		log.Fatal("Can't connect to database!")
	}

	db.Ping()

	//_, err := db.Db.Query("SELECT * FROM domains")
}
<commit_msg>a little test<commit_after>package ddns_server

import (
	log "github.com/Sirupsen/logrus"
	"github.com/meshwalker/ddns-server/master/database"
)

func main() {
	cfg, err := GetConfig()
	if err != nil {
		log.Fatal("Can't read config values")
	}

	db, err := database.New(cfg.DbHost, cfg.Port, cfg.DbName, cfg.DbUser, cfg.DbPassword)
	if err != nil {
		log.Fatal("Can't connect to database!")
	}

	db.Ping()

	//_, err := db.Db.Query("SELECT * FROM domains")
}
<|endoftext|>


<commit_before>package gracehttp

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

const (
	GRACEFUL_ENVIRON_KEY    = "IS_GRACEFUL"
	GRACEFUL_ENVIRON_STRING = GRACEFUL_ENVIRON_KEY + "=1"

	DEFAULT_READ_TIMEOUT  = 30 * time.Second
	DEFAULT_WRITE_TIMEOUT = DEFAULT_READ_TIMEOUT
)

// refer http.ListenAndServe
func ListenAndServe(addr string, handler http.Handler) error {
	return NewServer(addr, handler, DEFAULT_READ_TIMEOUT, DEFAULT_WRITE_TIMEOUT).ListenAndServe()
}

// refer http.ListenAndServeTLS
func ListenAndServeTLS(addr string, certFile string, keyFile string, handler http.Handler) error {
	return NewServer(addr, handler, DEFAULT_READ_TIMEOUT, DEFAULT_WRITE_TIMEOUT).ListenAndServeTLS(certFile, keyFile)
}

// new server
func NewServer(addr string, handler http.Handler, readTimeout, writeTimeout time.Duration) *Server {

	// 获取环境变量
	isGraceful := false
	if os.Getenv(GRACEFUL_ENVIRON_KEY) != "" {
		isGraceful = true
	}

	// 实例化Server
	return &Server{
		httpServer: &http.Server{
			Addr:    addr,
			Handler: handler,

			ReadTimeout:  readTimeout,
			WriteTimeout: writeTimeout,
		},

		isGraceful: isGraceful,
		signalChan: make(chan os.Signal),
	}
}

// 支持优雅重启的http服务
type Server struct {
	httpServer *http.Server
	listener   net.Listener

	isGraceful bool
	signalChan chan os.Signal
}

func (this *Server) ListenAndServe() error {
	addr := this.httpServer.Addr
	if addr == "" {
		addr = ":http"
	}

	ln, err := this.getNetTCPListener(addr)
	if err != nil {
		return err
	}

	this.listener = newListener(ln)

	return this.Serve()
}

func (this *Server) ListenAndServeTLS(certFile, keyFile string) error {
	addr := this.httpServer.Addr
	if addr == "" {
		addr = ":https"
	}

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

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

	ln, err := this.getNetTCPListener(addr)
	if err != nil {
		return err
	}

	this.listener = tls.NewListener(newListener(ln), config)
	return this.Serve()
}

func (this *Server) Serve() error {

	// 处理信号
	go this.handleSignals()

	// 处理HTTP请求
	err := this.httpServer.Serve(this.listener)

	// 跳出Serve处理代表 listener 已经close，等待所有已有的连接处理结束
	this.logf("waiting for connection close...")
	this.listener.(*Listener).Wait()
	this.logf("all connection closed, process with pid %d shutting down...", os.Getpid())

	return err
}

func (this *Server) getNetTCPListener(addr string) (*net.TCPListener, error) {

	var ln net.Listener
	var err error

	if this.isGraceful {
		file := os.NewFile(3, "")
		ln, err = net.FileListener(file)
		if err != nil {
			err = fmt.Errorf("net.FileListener error: %v", err)
			return nil, err
		}
	} else {
		ln, err = net.Listen("tcp", addr)
		if err != nil {
			err = fmt.Errorf("net.Listen error: %v", err)
			return nil, err
		}
	}
	return ln.(*net.TCPListener), nil
}

func (this *Server) handleSignals() {
	var sig os.Signal

	signal.Notify(
		this.signalChan,
		syscall.SIGHUP,
		syscall.SIGTERM,
	)

	pid := os.Getpid()
	for {
		sig = <-this.signalChan

		switch sig {

		case syscall.SIGTERM:

			this.logf("pid %d received SIGTERM.", pid)
			this.logf("graceful shutting down http server...")

			// 关闭老进程的连接
			this.listener.(*Listener).Close()
			this.logf("listener of pid %d closed.", pid)

		case syscall.SIGHUP:

			this.logf("pid %d received SIGHUP.", pid)
			this.logf("graceful restart http server...")

			err := this.startNewProcess()
			if err != nil {
				this.logf("start new process failed: %v, pid %d continue serve.", err)
			} else {
				// 关闭老进程的连接
				this.listener.(*Listener).Close()
				this.logf("listener of pid %d closed.", pid)
			}

		default:

		}
	}
}

// 启动子进程执行新程序
func (this *Server) startNewProcess() error {

	listenerFd, err := this.listener.(*Listener).GetFd()
	if err != nil {
		return fmt.Errorf("failed to get socket file descriptor: %v", err)
	}

	path := os.Args[0]

	// 设置标识优雅重启的环境变量
	environList := []string{}
	for _, value := range os.Environ() {
		if value != GRACEFUL_ENVIRON_STRING {
			environList = append(environList, value)
		}
	}
	environList = append(environList, GRACEFUL_ENVIRON_STRING)

	execSpec := &syscall.ProcAttr{
		Env:   environList,
		Files: []uintptr{os.Stdin.Fd(), os.Stdout.Fd(), os.Stderr.Fd(), listenerFd},
	}

	fork, err := syscall.ForkExec(path, os.Args, execSpec)
	if err != nil {
		return fmt.Errorf("failed to forkexec: %v", err)
	}

	this.logf("start new process success, pid %d.", fork)

	return nil
}

func (this *Server) logf(format string, args ...interface{}) {

	if this.httpServer.ErrorLog != nil {
		this.httpServer.ErrorLog.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}
<commit_msg>change SIGHUP to SIGUSR2<commit_after>package gracehttp

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

const (
	GRACEFUL_ENVIRON_KEY    = "IS_GRACEFUL"
	GRACEFUL_ENVIRON_STRING = GRACEFUL_ENVIRON_KEY + "=1"

	DEFAULT_READ_TIMEOUT  = 30 * time.Second
	DEFAULT_WRITE_TIMEOUT = DEFAULT_READ_TIMEOUT
)

// refer http.ListenAndServe
func ListenAndServe(addr string, handler http.Handler) error {
	return NewServer(addr, handler, DEFAULT_READ_TIMEOUT, DEFAULT_WRITE_TIMEOUT).ListenAndServe()
}

// refer http.ListenAndServeTLS
func ListenAndServeTLS(addr string, certFile string, keyFile string, handler http.Handler) error {
	return NewServer(addr, handler, DEFAULT_READ_TIMEOUT, DEFAULT_WRITE_TIMEOUT).ListenAndServeTLS(certFile, keyFile)
}

// new server
func NewServer(addr string, handler http.Handler, readTimeout, writeTimeout time.Duration) *Server {

	// 获取环境变量
	isGraceful := false
	if os.Getenv(GRACEFUL_ENVIRON_KEY) != "" {
		isGraceful = true
	}

	// 实例化Server
	return &Server{
		httpServer: &http.Server{
			Addr:    addr,
			Handler: handler,

			ReadTimeout:  readTimeout,
			WriteTimeout: writeTimeout,
		},

		isGraceful: isGraceful,
		signalChan: make(chan os.Signal),
	}
}

// 支持优雅重启的http服务
type Server struct {
	httpServer *http.Server
	listener   net.Listener

	isGraceful bool
	signalChan chan os.Signal
}

func (this *Server) ListenAndServe() error {
	addr := this.httpServer.Addr
	if addr == "" {
		addr = ":http"
	}

	ln, err := this.getNetTCPListener(addr)
	if err != nil {
		return err
	}

	this.listener = newListener(ln)

	return this.Serve()
}

func (this *Server) ListenAndServeTLS(certFile, keyFile string) error {
	addr := this.httpServer.Addr
	if addr == "" {
		addr = ":https"
	}

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

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

	ln, err := this.getNetTCPListener(addr)
	if err != nil {
		return err
	}

	this.listener = tls.NewListener(newListener(ln), config)
	return this.Serve()
}

func (this *Server) Serve() error {

	// 处理信号
	go this.handleSignals()

	// 处理HTTP请求
	err := this.httpServer.Serve(this.listener)

	// 跳出Serve处理代表 listener 已经close，等待所有已有的连接处理结束
	this.logf("waiting for connection close...")
	this.listener.(*Listener).Wait()
	this.logf("all connection closed, process with pid %d shutting down...", os.Getpid())

	return err
}

func (this *Server) getNetTCPListener(addr string) (*net.TCPListener, error) {

	var ln net.Listener
	var err error

	if this.isGraceful {
		file := os.NewFile(3, "")
		ln, err = net.FileListener(file)
		if err != nil {
			err = fmt.Errorf("net.FileListener error: %v", err)
			return nil, err
		}
	} else {
		ln, err = net.Listen("tcp", addr)
		if err != nil {
			err = fmt.Errorf("net.Listen error: %v", err)
			return nil, err
		}
	}
	return ln.(*net.TCPListener), nil
}

func (this *Server) handleSignals() {
	var sig os.Signal

	signal.Notify(
		this.signalChan,
		syscall.SIGTERM,
		syscall.SIGUSR2,
	)

	pid := os.Getpid()
	for {
		sig = <-this.signalChan

		switch sig {

		case syscall.SIGTERM:

			this.logf("pid %d received SIGTERM.", pid)
			this.logf("graceful shutting down http server...")

			// 关闭老进程的连接
			this.listener.(*Listener).Close()
			this.logf("listener of pid %d closed.", pid)

		case syscall.SIGUSR2:

			this.logf("pid %d received SIGUSR2.", pid)
			this.logf("graceful restart http server...")

			err := this.startNewProcess()
			if err != nil {
				this.logf("start new process failed: %v, pid %d continue serve.", err)
			} else {
				// 关闭老进程的连接
				this.listener.(*Listener).Close()
				this.logf("listener of pid %d closed.", pid)
			}

		default:

		}
	}
}

// 启动子进程执行新程序
func (this *Server) startNewProcess() error {

	listenerFd, err := this.listener.(*Listener).GetFd()
	if err != nil {
		return fmt.Errorf("failed to get socket file descriptor: %v", err)
	}

	path := os.Args[0]

	// 设置标识优雅重启的环境变量
	environList := []string{}
	for _, value := range os.Environ() {
		if value != GRACEFUL_ENVIRON_STRING {
			environList = append(environList, value)
		}
	}
	environList = append(environList, GRACEFUL_ENVIRON_STRING)

	execSpec := &syscall.ProcAttr{
		Env:   environList,
		Files: []uintptr{os.Stdin.Fd(), os.Stdout.Fd(), os.Stderr.Fd(), listenerFd},
	}

	fork, err := syscall.ForkExec(path, os.Args, execSpec)
	if err != nil {
		return fmt.Errorf("failed to forkexec: %v", err)
	}

	this.logf("start new process success, pid %d.", fork)

	return nil
}

func (this *Server) logf(format string, args ...interface{}) {

	if this.httpServer.ErrorLog != nil {
		this.httpServer.ErrorLog.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}
<|endoftext|>


<commit_before>package http

import (
	"github.com/ian-kent/go-log/log"
	gotcha "github.com/ian-kent/gotcha/app"
	"github.com/mailhog/MailHog-Server/config"
)

func Listen(cfg *config.Config, Asset func(string) ([]byte, error), exitCh chan int, registerCallback func(*gotcha.App)) {
	log.Info("[HTTP] Binding to address: %s", cfg.HTTPBindAddr)

	var app = gotcha.Create(Asset)
	app.Config.Listen = cfg.HTTPBindAddr

	registerCallback(app)

	app.Start()

	<-make(chan int)
	exitCh <- 1
}
<commit_msg>Remove dependency on config<commit_after>package http

import (
	"github.com/ian-kent/go-log/log"
	gotcha "github.com/ian-kent/gotcha/app"
)

func Listen(httpBindAddr string, Asset func(string) ([]byte, error), exitCh chan int, registerCallback func(*gotcha.App)) {
	log.Info("[HTTP] Binding to address: %s", httpBindAddr)

	var app = gotcha.Create(Asset)
	app.Config.Listen = httpBindAddr

	registerCallback(app)

	app.Start()

	<-make(chan int)
	exitCh <- 1
}
<|endoftext|>


<commit_before>package radius

import (
	"context"
	"errors"
	"net"
)

// ErrServerShutdown is returned from server Serve methods when Shutdown
// has been called and handlers are still completing.
var ErrServerShutdown = errors.New("radius: server is shutting down")

// Handler provides a handler to RADIUS server requests. When a RADIUS request
// is received, ServeRADIUS is called.
type Handler interface {
	ServeRADIUS(w ResponseWriter, r *Request)
}

// HandlerFunc allows a function to implement Handler.
type HandlerFunc func(w ResponseWriter, r *Request)

// ServeRADIUS calls h(w, p).
func (h HandlerFunc) ServeRADIUS(w ResponseWriter, r *Request) {
	h(w, r)
}

// Request is an incoming RADIUS request that is being handled by the server.
type Request struct {
	// LocalAddr is the local address on which the incoming RADIUS request
	// was received.
	LocalAddr net.Addr
	// RemoteAddr is the address from which the incoming RADIUS request
	// was sent.
	RemoteAddr net.Addr

	// Packet is the RADIUS packet sent in the request.
	*Packet

	ctx context.Context
}

// Context returns the context of the request. If a context has not been set
// using WithContext, the Background context is returned.
func (r *Request) Context() context.Context {
	if r.ctx != nil {
		return r.ctx
	}
	return context.Background()
}

// WithContext returns a shallow copy of the request with the new request's
// context set to the given context.
func (r *Request) WithContext(ctx context.Context) *Request {
	if ctx == nil {
		panic("nil ctx")
	}
	req := new(Request)
	*req = *r
	req.ctx = ctx
	return req
}

// ResponseWriter is used by RADIUS servers when replying to a RADIUS request.
type ResponseWriter interface {
	Write(packet *Packet) error
}

// SecretSource supplies RADIUS servers with the secret that should be used for
// authorizing and decrypting packets.
//
// ctx is canceled if the server's Shutdown method is called.
type SecretSource interface {
	RADIUSSecret(ctx context.Context, remoteAddr net.Addr) ([]byte, error)
}

// StaticSecretSource returns a SecretSource that uses secret for all requests.
func StaticSecretSource(secret []byte) SecretSource {
	return staticSecretSource(secret)
}

type staticSecretSource []byte

func (secret staticSecretSource) RADIUSSecret(ctx context.Context, remoteAddr net.Addr) ([]byte, error) {
	return []byte(secret), nil
}
<commit_msg>wrap staticSecretSource secret in struct<commit_after>package radius

import (
	"context"
	"errors"
	"net"
)

// ErrServerShutdown is returned from server Serve methods when Shutdown
// has been called and handlers are still completing.
var ErrServerShutdown = errors.New("radius: server is shutting down")

// Handler provides a handler to RADIUS server requests. When a RADIUS request
// is received, ServeRADIUS is called.
type Handler interface {
	ServeRADIUS(w ResponseWriter, r *Request)
}

// HandlerFunc allows a function to implement Handler.
type HandlerFunc func(w ResponseWriter, r *Request)

// ServeRADIUS calls h(w, p).
func (h HandlerFunc) ServeRADIUS(w ResponseWriter, r *Request) {
	h(w, r)
}

// Request is an incoming RADIUS request that is being handled by the server.
type Request struct {
	// LocalAddr is the local address on which the incoming RADIUS request
	// was received.
	LocalAddr net.Addr
	// RemoteAddr is the address from which the incoming RADIUS request
	// was sent.
	RemoteAddr net.Addr

	// Packet is the RADIUS packet sent in the request.
	*Packet

	ctx context.Context
}

// Context returns the context of the request. If a context has not been set
// using WithContext, the Background context is returned.
func (r *Request) Context() context.Context {
	if r.ctx != nil {
		return r.ctx
	}
	return context.Background()
}

// WithContext returns a shallow copy of the request with the new request's
// context set to the given context.
func (r *Request) WithContext(ctx context.Context) *Request {
	if ctx == nil {
		panic("nil ctx")
	}
	req := new(Request)
	*req = *r
	req.ctx = ctx
	return req
}

// ResponseWriter is used by RADIUS servers when replying to a RADIUS request.
type ResponseWriter interface {
	Write(packet *Packet) error
}

// SecretSource supplies RADIUS servers with the secret that should be used for
// authorizing and decrypting packets.
//
// ctx is canceled if the server's Shutdown method is called.
type SecretSource interface {
	RADIUSSecret(ctx context.Context, remoteAddr net.Addr) ([]byte, error)
}

// StaticSecretSource returns a SecretSource that uses secret for all requests.
func StaticSecretSource(secret []byte) SecretSource {
	return &staticSecretSource{secret}
}

type staticSecretSource struct {
	secret []byte
}

func (s *staticSecretSource) RADIUSSecret(ctx context.Context, remoteAddr net.Addr) ([]byte, error) {
	return s.secret, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/deferpanic/deferclient/deferstats"
	"github.com/fzzy/radix/extra/pool"
	"github.com/fzzy/radix/redis"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/gorilla/securecookie"
	"github.com/namsral/flag"
	"github.com/unrolled/render"
	// "github.com/deferpanic/deferclient/errors"
	"golang.org/x/oauth2"
	githuboauth "golang.org/x/oauth2/github"
	"log"
	"math/rand"
	"net/http"
	"runtime"
	"strings"
	"time"
)

// func QueryScore(terms []string, title) float32 {
// 	return 1.0
// }

func errorHandler(err error) {
	if err != nil {
		log.Panic(err)
	}
}

func isOnHTTPS(req *http.Request) bool {
	if req.Header.Get("is-secure") == "true" {
		return true
	}
	// default is to use the flag
	// which is only really useful for local development
	return usingHTTPS
}

type domainRow struct {
	Key    string
	Domain string
}

func getUsername(req *http.Request) string {
	var username string
	cookie, err := req.Cookie("username")
	if err == nil {
		sCookie.Decode("username", cookie.Value, &username)
	}
	return username
}

func indexHandler(w http.ResponseWriter, req *http.Request) {
	context := map[string]interface{}{
		"staticPrefix": staticPrefix,
		"isNotDebug":   !debug,
		"Username":     "",
		"domains":      make([]string, 0),
	}

	username := getUsername(req)
	if username != "" {
		// Yay! You're signed in!
		context["Username"] = username
		c, err := redisPool.Get()
		errorHandler(err)
		defer redisPool.CarefullyPut(c, &err)

		userdomainsKey := fmt.Sprintf("$userdomains$%v", username)
		replies, err := c.Cmd("SMEMBERS", userdomainsKey).List()
		errorHandler(err)

		var domains []domainRow

		var domain string
		for _, key := range replies {
			reply := c.Cmd("HGET", "$domainkeys", key)
			if reply.Type != redis.NilReply {
				domain, err = reply.Str()
				errorHandler(err)
				domains = append(domains, domainRow{
					Key:    key,
					Domain: domain,
				})
			}
		}
		context["domains"] = domains
	}
	// this assumes there's a `templates/index.tmpl` file
	renderer.HTML(w, http.StatusOK, "index", context)
}

func logoutHandler(w http.ResponseWriter, req *http.Request) {
	expire := time.Now().AddDate(0, 0, -1)
	secureCookie := isOnHTTPS(req)
	cookie := &http.Cookie{
		Name:     "username",
		Value:    "*deleted*",
		Path:     "/",
		Expires:  expire,
		MaxAge:   -1,
		Secure:   secureCookie,
		HttpOnly: true,
	}
	http.SetCookie(w, cookie)
	http.Redirect(w, req, "/#loggedout", http.StatusTemporaryRedirect)
}

func handleGitHubLogin(w http.ResponseWriter, req *http.Request) {
	url := oauthConf.AuthCodeURL(oauthStateString, oauth2.AccessTypeOnline)
	http.Redirect(w, req, url, http.StatusTemporaryRedirect)
}

func handleGitHubCallback(w http.ResponseWriter, req *http.Request) {
	state := req.FormValue("state")
	if state != oauthStateString {
		log.Printf(
			"invalid oauth state, expected '%s', got '%s'\n",
			oauthStateString,
			state,
		)
		http.Redirect(w, req, "/", http.StatusTemporaryRedirect)
		return
	}

	code := req.FormValue("code")
	token, err := oauthConf.Exchange(oauth2.NoContext, code)
	if err != nil {
		log.Printf("oauthConf.Exchange() failed with '%s'\n", err)
		http.Redirect(w, req, "/", http.StatusTemporaryRedirect)
		return
	}

	oauthClient := oauthConf.Client(oauth2.NoContext, token)
	client := github.NewClient(oauthClient)
	// the second item here is the github.Rate config
	user, _, err := client.Users.Get("")

	if err != nil {
		log.Printf("client.Users.Get() faled with '%s'\n", err)
		http.Redirect(w, req, "/", http.StatusTemporaryRedirect)
		return
	}

	log.Printf("Logged in as GitHub user: %s (%s)\n", *user.Login, *user.Email)
	// fmt.Println(user)
	// store the email
	c, err := redisPool.Get()
	errorHandler(err)
	defer redisPool.CarefullyPut(c, &err)

	err = c.Cmd("SADD", "$usernames", *user.Login).Err
	errorHandler(err)
	err = c.Cmd("HSET", "$emails", *user.Login, *user.Email).Err
	errorHandler(err)

	encoded, err := sCookie.Encode("username", *user.Login)
	errorHandler(err)
	expire := time.Now().AddDate(0, 0, 1) // how long is this?
	secureCookie := isOnHTTPS(req)

	cookie := &http.Cookie{
		Name:     "username",
		Value:    encoded,
		Path:     "/",
		Expires:  expire,
		MaxAge:   60 * 60 * 24 * 30, // 30 days
		Secure:   secureCookie,
		HttpOnly: true,
	}
	http.SetCookie(w, cookie)
	http.Redirect(w, req, "/#auth", http.StatusTemporaryRedirect)
}

var letters = []rune(
	"abcdefghjkmnopqrstuvwxyzABCDEFGHJKMNPQRSTUVWXYZ123456789",
)

func randString(n int) string {
	rand.Seed(time.Now().UTC().UnixNano())
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

func domainkeyNewHandler(w http.ResponseWriter, req *http.Request) {
	domain := strings.Trim(req.FormValue("domain"), " ")
	if domain != "" {
		username := getUsername(req)
		if username != "" {
			c, err := redisPool.Get()
			errorHandler(err)
			defer redisPool.CarefullyPut(c, &err)

			key := randString(24)
			userdomainsKey := fmt.Sprintf("$userdomains$%v", username)
			err = c.Cmd("SADD", userdomainsKey, key).Err
			errorHandler(err)
			err = c.Cmd("HSET", "$domainkeys", key, domain).Err
			errorHandler(err)
		}
	}

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

func domainkeyDeleteHandler(w http.ResponseWriter, req *http.Request) {
	key := strings.Trim(req.FormValue("key"), " ")
	if key != "" {
		username := getUsername(req)
		if username != "" {
			// Yay! You're signed in!
			c, err := redisPool.Get()
			errorHandler(err)
			defer redisPool.CarefullyPut(c, &err)

			userdomainsKey := fmt.Sprintf("$userdomains$%v", username)
			err = c.Cmd("SREM", userdomainsKey, key).Err
			errorHandler(err)
			err = c.Cmd("HDEL", "$domainkeys", key).Err
			errorHandler(err)
		} // else, we should yield some sort of 403 message maybe
	}
	http.Redirect(w, req, "/#auth", http.StatusFound)
}

var (
	redisPool    *pool.Pool
	procs        int
	debug        = true
	renderer     = render.New()
	redisURL     = "127.0.0.1:6379"
	staticPrefix = ""
	usingHTTPS   = false
	sCookie      *securecookie.SecureCookie
)

var (
	// You must register the app at https://github.com/settings/applications
	// Set callback to http://127.0.0.1:7000/github_oauth_cb
	// Set ClientId and ClientSecret to
	oauthConf = &oauth2.Config{
		ClientID:     "",
		ClientSecret: "",
		Scopes:       []string{"user:email"},
		Endpoint:     githuboauth.Endpoint,
	}
	// random string for oauth2 API calls to protect against CSRF
	oauthStateString = randString(24)
)

func main() {
	var (
		port          = 3001
		redisDatabase = 0
		redisPoolSize = 10
		clientID      = ""
		clientSecret  = ""
		hashKey       = "randomishstringthatsi32charslong"
		blockKey      = "randomishstringthatsi32charslong"
		deferPanicKey = ""
	)
	flag.IntVar(&port, "port", port, "Port to start the server on")
	flag.IntVar(&procs, "procs", 1, "Number of CPU processors (0 to use max)")
	flag.BoolVar(&debug, "debug", false, "Debug mode")
	flag.StringVar(
		&redisURL, "redisURL", redisURL,
		"Redis URL to tcp connect to")
	flag.StringVar(
		&staticPrefix, "staticPrefix", staticPrefix,
		"Prefix in front of static assets in HTML")
	flag.IntVar(&redisDatabase, "redisDatabase", redisDatabase,
		"Redis database number to connect to")
	flag.StringVar(
		&clientID, "clientID", clientID,
		"OAuth Client ID")
	flag.StringVar(
		&clientSecret, "clientSecret", clientSecret,
		"OAuth Client Secret")
	flag.BoolVar(&usingHTTPS, "usingHTTPS", usingHTTPS,
		"Whether requests are made under HTTPS")
	flag.StringVar(
		&hashKey, "hashKey", hashKey,
		"HMAC hash key to use for encoding cookies")
	flag.StringVar(
		&blockKey, "blockKey", blockKey,
		"Block key to encrypt cookie values")
	flag.StringVar(
		&deferPanicKey, "deferPanicKey", deferPanicKey,
		"Auth key for deferpanic.com")
	flag.Parse()

	dfs := deferstats.NewClient(deferPanicKey)
	go dfs.CaptureStats()
	// if deferPanicKey != "" {
	// 	dfs := deferstats.NewClient(deferPanicKey)
	// 	go dfs.CaptureStats()
	// }

	oauthConf.ClientID = clientID
	oauthConf.ClientSecret = clientSecret

	sCookie = securecookie.New([]byte(hashKey), []byte(blockKey))

	log.Println("REDIS DATABASE:", redisDatabase)
	log.Println("DEBUG MODE:", debug)
	log.Println("STATIC PREFIX:", staticPrefix)

	if !debug {
		redisPoolSize = 100
	}

	// Figuring out how many processors to use.
	maxProcs := runtime.NumCPU()
	if procs == 0 {
		procs = maxProcs
	} else if procs < 0 {
		panic("PROCS < 0")
	} else if procs > maxProcs {
		panic(fmt.Sprintf("PROCS > max (%v)", maxProcs))
	}
	log.Println("PROCS:", procs)
	runtime.GOMAXPROCS(procs)

	renderer = render.New(render.Options{
		IndentJSON:    debug,
		IsDevelopment: debug,
	})

	df := func(network, addr string) (*redis.Client, error) {
		client, err := redis.Dial(network, addr)
		if err != nil {
			return nil, err
		}
		err = client.Cmd("SELECT", redisDatabase).Err
		if err != nil {
			return nil, err
		}
		// if err = client.Cmd("AUTH", "SUPERSECRET").Err; err != nil {
		// 	client.Close()
		// 	return nil, err
		// }
		return client, nil
	}

	var err error
	redisPool, err = pool.NewCustomPool("tcp", redisURL, redisPoolSize, df)
	errorHandler(err)

	mux := mux.NewRouter()
	mux.HandleFunc("/", dfs.HTTPHandler(indexHandler)).Methods("GET", "HEAD")
	mux.HandleFunc("/v1/ping", dfs.HTTPHandler(pingHandler)).Methods("GET", "HEAD")
	mux.HandleFunc("/v1", dfs.HTTPHandler(fetchHandler)).Methods("GET", "HEAD")
	mux.HandleFunc("/v1", dfs.HTTPHandler(updateHandler)).Methods("POST", "PUT")
	mux.HandleFunc("/v1", dfs.HTTPHandler(deleteHandler)).Methods("DELETE")
	mux.HandleFunc("/v1/stats", dfs.HTTPHandler(privateStatsHandler)).Methods("GET")
	mux.HandleFunc("/v1/flush", dfs.HTTPHandler(flushHandler)).Methods("DELETE")
	mux.HandleFunc("/v1/bulk", dfs.HTTPHandler(bulkHandler)).Methods("POST", "PUT")
	mux.HandleFunc("/login", dfs.HTTPHandler(handleGitHubLogin)).Methods("GET")
	mux.HandleFunc("/logout", dfs.HTTPHandler(logoutHandler)).Methods("GET", "POST")
	mux.HandleFunc("/github_oauth_cb", dfs.HTTPHandler(handleGitHubCallback)).Methods("GET")
	mux.HandleFunc("/domainkeys/new", dfs.HTTPHandler(domainkeyNewHandler)).Methods("POST")
	mux.HandleFunc("/domainkeys/delete", dfs.HTTPHandler(domainkeyDeleteHandler)).Methods("POST")

	n := negroni.Classic()

	n.UseHandler(mux)
	n.Run(fmt.Sprintf(":%d", port))
}
<commit_msg>more careful about github email<commit_after>package main

import (
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/deferpanic/deferclient/deferstats"
	"github.com/fzzy/radix/extra/pool"
	"github.com/fzzy/radix/redis"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/gorilla/securecookie"
	"github.com/namsral/flag"
	"github.com/unrolled/render"
	// "github.com/deferpanic/deferclient/errors"
	"golang.org/x/oauth2"
	githuboauth "golang.org/x/oauth2/github"
	"log"
	"math/rand"
	"net/http"
	"runtime"
	"strings"
	"time"
)

// func QueryScore(terms []string, title) float32 {
// 	return 1.0
// }

func errorHandler(err error) {
	if err != nil {
		log.Panic(err)
	}
}

func isOnHTTPS(req *http.Request) bool {
	if req.Header.Get("is-secure") == "true" {
		return true
	}
	// default is to use the flag
	// which is only really useful for local development
	return usingHTTPS
}

type domainRow struct {
	Key    string
	Domain string
}

func getUsername(req *http.Request) string {
	var username string
	cookie, err := req.Cookie("username")
	if err == nil {
		sCookie.Decode("username", cookie.Value, &username)
	}
	return username
}

func indexHandler(w http.ResponseWriter, req *http.Request) {
	context := map[string]interface{}{
		"staticPrefix": staticPrefix,
		"isNotDebug":   !debug,
		"Username":     "",
		"domains":      make([]string, 0),
	}

	username := getUsername(req)
	if username != "" {
		// Yay! You're signed in!
		context["Username"] = username
		c, err := redisPool.Get()
		errorHandler(err)
		defer redisPool.CarefullyPut(c, &err)

		userdomainsKey := fmt.Sprintf("$userdomains$%v", username)
		replies, err := c.Cmd("SMEMBERS", userdomainsKey).List()
		errorHandler(err)

		var domains []domainRow

		var domain string
		for _, key := range replies {
			reply := c.Cmd("HGET", "$domainkeys", key)
			if reply.Type != redis.NilReply {
				domain, err = reply.Str()
				errorHandler(err)
				domains = append(domains, domainRow{
					Key:    key,
					Domain: domain,
				})
			}
		}
		context["domains"] = domains
	}
	// this assumes there's a `templates/index.tmpl` file
	renderer.HTML(w, http.StatusOK, "index", context)
}

func logoutHandler(w http.ResponseWriter, req *http.Request) {
	expire := time.Now().AddDate(0, 0, -1)
	secureCookie := isOnHTTPS(req)
	cookie := &http.Cookie{
		Name:     "username",
		Value:    "*deleted*",
		Path:     "/",
		Expires:  expire,
		MaxAge:   -1,
		Secure:   secureCookie,
		HttpOnly: true,
	}
	http.SetCookie(w, cookie)
	http.Redirect(w, req, "/#loggedout", http.StatusTemporaryRedirect)
}

func handleGitHubLogin(w http.ResponseWriter, req *http.Request) {
	url := oauthConf.AuthCodeURL(oauthStateString, oauth2.AccessTypeOnline)
	http.Redirect(w, req, url, http.StatusTemporaryRedirect)
}

func handleGitHubCallback(w http.ResponseWriter, req *http.Request) {
	state := req.FormValue("state")
	if state != oauthStateString {
		log.Printf(
			"invalid oauth state, expected '%s', got '%s'\n",
			oauthStateString,
			state,
		)
		http.Redirect(w, req, "/", http.StatusTemporaryRedirect)
		return
	}

	code := req.FormValue("code")
	token, err := oauthConf.Exchange(oauth2.NoContext, code)
	if err != nil {
		log.Printf("oauthConf.Exchange() failed with '%s'\n", err)
		http.Redirect(w, req, "/", http.StatusTemporaryRedirect)
		return
	}

	oauthClient := oauthConf.Client(oauth2.NoContext, token)
	client := github.NewClient(oauthClient)
	// the second item here is the github.Rate config
	user, _, err := client.Users.Get("")

	if err != nil {
		log.Printf("client.Users.Get() faled with '%s'\n", err)
		http.Redirect(w, req, "/", http.StatusTemporaryRedirect)
		return
	}

	log.Printf("Logged in as GitHub user: %s (%s)\n", *user.Login, *user.Email)
	// fmt.Println(user)
	// store the email
	c, err := redisPool.Get()
	errorHandler(err)
	defer redisPool.CarefullyPut(c, &err)

	err = c.Cmd("SADD", "$usernames", *user.Login).Err
	errorHandler(err)
	if user.Email != nil {
		err = c.Cmd("HSET", "$emails", *user.Login, *user.Email).Err
		errorHandler(err)
	}
	// c.Cmd("HSET", "$emails", *user.Login, *user.Email)
	// fmt.Println(*user.Login, *user.Email)

	encoded, err := sCookie.Encode("username", *user.Login)
	errorHandler(err)
	expire := time.Now().AddDate(0, 0, 1) // how long is this?
	secureCookie := isOnHTTPS(req)

	cookie := &http.Cookie{
		Name:     "username",
		Value:    encoded,
		Path:     "/",
		Expires:  expire,
		MaxAge:   60 * 60 * 24 * 30, // 30 days
		Secure:   secureCookie,
		HttpOnly: true,
	}
	http.SetCookie(w, cookie)
	http.Redirect(w, req, "/#auth", http.StatusTemporaryRedirect)
}

var letters = []rune(
	"abcdefghjkmnopqrstuvwxyzABCDEFGHJKMNPQRSTUVWXYZ123456789",
)

func randString(n int) string {
	rand.Seed(time.Now().UTC().UnixNano())
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

func domainkeyNewHandler(w http.ResponseWriter, req *http.Request) {
	domain := strings.Trim(req.FormValue("domain"), " ")
	if domain != "" {
		username := getUsername(req)
		if username != "" {
			c, err := redisPool.Get()
			errorHandler(err)
			defer redisPool.CarefullyPut(c, &err)

			key := randString(24)
			userdomainsKey := fmt.Sprintf("$userdomains$%v", username)
			err = c.Cmd("SADD", userdomainsKey, key).Err
			errorHandler(err)
			err = c.Cmd("HSET", "$domainkeys", key, domain).Err
			errorHandler(err)
		}
	}

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

func domainkeyDeleteHandler(w http.ResponseWriter, req *http.Request) {
	key := strings.Trim(req.FormValue("key"), " ")
	if key != "" {
		username := getUsername(req)
		if username != "" {
			// Yay! You're signed in!
			c, err := redisPool.Get()
			errorHandler(err)
			defer redisPool.CarefullyPut(c, &err)

			userdomainsKey := fmt.Sprintf("$userdomains$%v", username)
			err = c.Cmd("SREM", userdomainsKey, key).Err
			errorHandler(err)
			err = c.Cmd("HDEL", "$domainkeys", key).Err
			errorHandler(err)
		} // else, we should yield some sort of 403 message maybe
	}
	http.Redirect(w, req, "/#auth", http.StatusFound)
}

var (
	redisPool    *pool.Pool
	procs        int
	debug        = true
	renderer     = render.New()
	redisURL     = "127.0.0.1:6379"
	staticPrefix = ""
	usingHTTPS   = false
	sCookie      *securecookie.SecureCookie
)

var (
	// You must register the app at https://github.com/settings/applications
	// Set callback to http://127.0.0.1:7000/github_oauth_cb
	// Set ClientId and ClientSecret to
	oauthConf = &oauth2.Config{
		ClientID:     "",
		ClientSecret: "",
		Scopes:       []string{"user:email"},
		Endpoint:     githuboauth.Endpoint,
	}
	// random string for oauth2 API calls to protect against CSRF
	oauthStateString = randString(24)
)

func main() {
	var (
		port          = 3001
		redisDatabase = 0
		redisPoolSize = 10
		clientID      = ""
		clientSecret  = ""
		hashKey       = "randomishstringthatsi32charslong"
		blockKey      = "randomishstringthatsi32charslong"
		deferPanicKey = ""
	)
	flag.IntVar(&port, "port", port, "Port to start the server on")
	flag.IntVar(&procs, "procs", 1, "Number of CPU processors (0 to use max)")
	flag.BoolVar(&debug, "debug", false, "Debug mode")
	flag.StringVar(
		&redisURL, "redisURL", redisURL,
		"Redis URL to tcp connect to")
	flag.StringVar(
		&staticPrefix, "staticPrefix", staticPrefix,
		"Prefix in front of static assets in HTML")
	flag.IntVar(&redisDatabase, "redisDatabase", redisDatabase,
		"Redis database number to connect to")
	flag.StringVar(
		&clientID, "clientID", clientID,
		"OAuth Client ID")
	flag.StringVar(
		&clientSecret, "clientSecret", clientSecret,
		"OAuth Client Secret")
	flag.BoolVar(&usingHTTPS, "usingHTTPS", usingHTTPS,
		"Whether requests are made under HTTPS")
	flag.StringVar(
		&hashKey, "hashKey", hashKey,
		"HMAC hash key to use for encoding cookies")
	flag.StringVar(
		&blockKey, "blockKey", blockKey,
		"Block key to encrypt cookie values")
	flag.StringVar(
		&deferPanicKey, "deferPanicKey", deferPanicKey,
		"Auth key for deferpanic.com")
	flag.Parse()

	dfs := deferstats.NewClient(deferPanicKey)
	go dfs.CaptureStats()
	// if deferPanicKey != "" {
	// 	dfs := deferstats.NewClient(deferPanicKey)
	// 	go dfs.CaptureStats()
	// }

	oauthConf.ClientID = clientID
	oauthConf.ClientSecret = clientSecret

	sCookie = securecookie.New([]byte(hashKey), []byte(blockKey))

	log.Println("REDIS DATABASE:", redisDatabase)
	log.Println("DEBUG MODE:", debug)
	log.Println("STATIC PREFIX:", staticPrefix)

	if !debug {
		redisPoolSize = 100
	}

	// Figuring out how many processors to use.
	maxProcs := runtime.NumCPU()
	if procs == 0 {
		procs = maxProcs
	} else if procs < 0 {
		panic("PROCS < 0")
	} else if procs > maxProcs {
		panic(fmt.Sprintf("PROCS > max (%v)", maxProcs))
	}
	log.Println("PROCS:", procs)
	runtime.GOMAXPROCS(procs)

	renderer = render.New(render.Options{
		IndentJSON:    debug,
		IsDevelopment: debug,
	})

	df := func(network, addr string) (*redis.Client, error) {
		client, err := redis.Dial(network, addr)
		if err != nil {
			return nil, err
		}
		err = client.Cmd("SELECT", redisDatabase).Err
		if err != nil {
			return nil, err
		}
		// if err = client.Cmd("AUTH", "SUPERSECRET").Err; err != nil {
		// 	client.Close()
		// 	return nil, err
		// }
		return client, nil
	}

	var err error
	redisPool, err = pool.NewCustomPool("tcp", redisURL, redisPoolSize, df)
	errorHandler(err)

	mux := mux.NewRouter()
	mux.HandleFunc("/", dfs.HTTPHandler(indexHandler)).Methods("GET", "HEAD")
	mux.HandleFunc("/v1/ping", dfs.HTTPHandler(pingHandler)).Methods("GET", "HEAD")
	mux.HandleFunc("/v1", dfs.HTTPHandler(fetchHandler)).Methods("GET", "HEAD")
	mux.HandleFunc("/v1", dfs.HTTPHandler(updateHandler)).Methods("POST", "PUT")
	mux.HandleFunc("/v1", dfs.HTTPHandler(deleteHandler)).Methods("DELETE")
	mux.HandleFunc("/v1/stats", dfs.HTTPHandler(privateStatsHandler)).Methods("GET")
	mux.HandleFunc("/v1/flush", dfs.HTTPHandler(flushHandler)).Methods("DELETE")
	mux.HandleFunc("/v1/bulk", dfs.HTTPHandler(bulkHandler)).Methods("POST", "PUT")
	mux.HandleFunc("/login", dfs.HTTPHandler(handleGitHubLogin)).Methods("GET")
	mux.HandleFunc("/logout", dfs.HTTPHandler(logoutHandler)).Methods("GET", "POST")
	mux.HandleFunc("/github_oauth_cb", dfs.HTTPHandler(handleGitHubCallback)).Methods("GET")
	mux.HandleFunc("/domainkeys/new", dfs.HTTPHandler(domainkeyNewHandler)).Methods("POST")
	mux.HandleFunc("/domainkeys/delete", dfs.HTTPHandler(domainkeyDeleteHandler)).Methods("POST")

	n := negroni.Classic()

	n.UseHandler(mux)
	n.Run(fmt.Sprintf(":%d", port))
}
<|endoftext|>


<commit_before>// Tideland Go REST Server Library - JSON Web Token - Unit Tests
//
// Copyright (C) 2016-2017 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package jwt_test

//--------------------
// IMPORTS
//--------------------

import (
	"context"
	"net/http"
	"testing"
	"time"

	"github.com/tideland/golib/audit"
	"github.com/tideland/golib/etc"

	"github.com/tideland/gorest/jwt"
	"github.com/tideland/gorest/rest"
	"github.com/tideland/gorest/restaudit"
)

//--------------------
// TESTS
//--------------------

// TestDecodeInvalidRequest tests the decoding of requests
// without a header or an invalid one.
func TestDecodeInvalidRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode invalid requests")
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	err := mux.Register("test", "jwt", NewTestHandler("jwt", assert, nil, false))
	assert.Nil(err)
	// Perform request without authorization.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	ts.DoRequest(req)
}

// TestDecodeRequest tests the decoding of a token
// in a handler.
func TestDecodeRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode a request token")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	err = mux.Register("test", "jwt", NewTestHandler("jwt", assert, nil, false))
	assert.Nil(err)
	// Perform test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestDecodeCachedRequest tests the decoding of a token
// in a handler including usage of the cache.
func TestDecodeCachedRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode a request token using a cache")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	err = mux.Register("test", "jwt", NewTestHandler("jwt", assert, nil, true))
	assert.Nil(err)
	// Perform first test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
	// Perform second test request.
	resp = ts.DoRequest(req)
	claimsOut = jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestVerifyRequest tests the verification of a token
// in a handler.
func TestVerifyRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing verify a request token")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	err = mux.Register("test", "jwt", NewTestHandler("jwt", assert, key, false))
	assert.Nil(err)
	// Perform test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestVerifyCachedRequest tests the verification of a token
// in a handler including usage of the cache.
func TestVerifyCachedRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing verify a request token using a cache")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	err = mux.Register("test", "jwt", NewTestHandler("jwt", assert, key, true))
	assert.Nil(err)
	// Perform first test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
	// Perform second test request.
	resp = ts.DoRequest(req)
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

//--------------------
// HANDLER
//--------------------

// testHandler is used in the test scenarios.
type testHandler struct {
	id     string
	assert audit.Assertion
	key    jwt.Key
	cache  jwt.Cache
}

func NewTestHandler(id string, assert audit.Assertion, key jwt.Key, useCache bool) rest.ResourceHandler {
	var cache jwt.Cache
	if useCache {
		cache = jwt.NewCache(time.Minute, time.Minute, time.Minute, 10)
	}
	return &testHandler{id, assert, key, cache}
}

func (th *testHandler) ID() string {
	return th.id
}

func (th *testHandler) Init(env rest.Environment, domain, resource string) error {
	return nil
}

func (th *testHandler) Get(job rest.Job) (bool, error) {
	if th.key == nil {
		return th.testDecode(job)
	} else {
		return th.testVerify(job)
	}
}

func (th *testHandler) testDecode(job rest.Job) (bool, error) {
	decode := func() (jwt.JWT, error) {
		if th.cache == nil {
			return jwt.DecodeFromJob(job)
		}
		return jwt.DecodeCachedFromJob(job, th.cache)
	}
	jwtOut, err := decode()
	th.assert.Nil(err)
	th.assert.True(jwtOut.IsValid(time.Minute))
	subject, ok := jwtOut.Claims().Subject()
	th.assert.True(ok)
	th.assert.Equal(subject, job.ResourceID())
	job.JSON(true).Write(rest.StatusOK, jwtOut.Claims())
	return true, nil
}

func (th *testHandler) testVerify(job rest.Job) (bool, error) {
	verify := func() (jwt.JWT, error) {
		if th.cache == nil {
			return jwt.VerifyFromJob(job, th.key)
		}
		return jwt.VerifyCachedFromJob(job, th.cache, th.key)
	}
	jwtOut, err := verify()
	th.assert.Nil(err)
	th.assert.True(jwtOut.IsValid(time.Minute))
	subject, ok := jwtOut.Claims().Subject()
	th.assert.True(ok)
	th.assert.Equal(subject, job.ResourceID())
	job.JSON(true).Write(rest.StatusOK, jwtOut.Claims())
	return true, nil
}

//--------------------
// HELPERS
//--------------------

// newMultiplexer creates a new multiplexer with a testing context
// and a testing configuration.
func newMultiplexer(assert audit.Assertion) rest.Multiplexer {
	cfgStr := "{etc {basepath /}{default-domain default}{default-resource default}}"
	cfg, err := etc.ReadString(cfgStr)
	assert.Nil(err)
	return rest.NewMultiplexer(context.Background(), cfg)
}

// EOF
<commit_msg>Redesigned JWT test handler, invalid tests can follow<commit_after>// Tideland Go REST Server Library - JSON Web Token - Unit Tests
//
// Copyright (C) 2016-2017 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package jwt_test

//--------------------
// IMPORTS
//--------------------

import (
	"context"
	"net/http"
	"testing"
	"time"

	"github.com/tideland/golib/audit"
	"github.com/tideland/golib/etc"

	"github.com/tideland/gorest/jwt"
	"github.com/tideland/gorest/rest"
	"github.com/tideland/gorest/restaudit"
)

//--------------------
// TESTS
//--------------------

// TestDecodeInvalidRequest tests the decoding of requests
// without a header or an invalid one.
func TestDecodeInvalidRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode invalid requests")
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newDecodeAsserter(assert, false)
	err := mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform request without authorization.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	ts.DoRequest(req)
}

// TestDecodeRequest tests the decoding of a token
// in a handler.
func TestDecodeRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode a request token")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newDecodeAsserter(assert, false)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestDecodeCachedRequest tests the decoding of a token
// in a handler including usage of the cache.
func TestDecodeCachedRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode a request token using a cache")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newDecodeAsserter(assert, true)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform first test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
	// Perform second test request.
	resp = ts.DoRequest(req)
	claimsOut = jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestVerifyRequest tests the verification of a token
// in a handler.
func TestVerifyRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing verify a request token")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newVerifyAsserter(assert, key, false)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestVerifyCachedRequest tests the verification of a token
// in a handler including usage of the cache.
func TestVerifyCachedRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing verify a request token using a cache")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newVerifyAsserter(assert, key, true)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform first test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
	// Perform second test request.
	resp = ts.DoRequest(req)
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

//--------------------
// HANDLER
//--------------------

// testAsserter instances will handle the assertions in the testHandler.
type testAsserter func(job rest.Job) (bool, error)

func newDecodeAsserter(assert audit.Assertion, cached bool) testAsserter {
	var cache jwt.Cache
	if cached {
		cache = jwt.NewCache(time.Minute, time.Minute, time.Minute, 10)
	}
	return func(job rest.Job) (bool, error) {
		var token jwt.JWT
		var err error
		if cached {
			token, err = jwt.DecodeCachedFromJob(job, cache)
		} else {
			token, err = jwt.DecodeFromJob(job)
		}
		assert.Nil(err)
		assert.True(token.IsValid(time.Minute))
		subject, ok := token.Claims().Subject()
		assert.True(ok)
		assert.Equal(subject, job.ResourceID())
		job.JSON(true).Write(rest.StatusOK, token.Claims())
		return true, nil
	}
}

func newVerifyAsserter(assert audit.Assertion, key jwt.Key, cached bool) testAsserter {
	var cache jwt.Cache
	if cached {
		cache = jwt.NewCache(time.Minute, time.Minute, time.Minute, 10)
	}
	return func(job rest.Job) (bool, error) {
		var token jwt.JWT
		var err error
		if cached {
			token, err = jwt.VerifyCachedFromJob(job, cache, key)
		} else {
			token, err = jwt.VerifyFromJob(job, key)
		}
		assert.Nil(err)
		assert.True(token.IsValid(time.Minute))
		subject, ok := token.Claims().Subject()
		assert.True(ok)
		assert.Equal(subject, job.ResourceID())
		job.JSON(true).Write(rest.StatusOK, token.Claims())
		return true, nil
	}
}

// testHandler is used in the test scenarios.
type testHandler struct {
	id       string
	asserter testAsserter
}

func newTestHandler(id string, asserter testAsserter) rest.ResourceHandler {
	return &testHandler{
		id:       id,
		asserter: asserter,
	}
}

func (th *testHandler) ID() string {
	return th.id
}

func (th *testHandler) Init(env rest.Environment, domain, resource string) error {
	return nil
}

func (th *testHandler) Get(job rest.Job) (bool, error) {
	return th.asserter(job)
}

//--------------------
// HELPERS
//--------------------

// newMultiplexer creates a new multiplexer with a testing context
// and a testing configuration.
func newMultiplexer(assert audit.Assertion) rest.Multiplexer {
	cfgStr := "{etc {basepath /}{default-domain default}{default-resource default}}"
	cfg, err := etc.ReadString(cfgStr)
	assert.Nil(err)
	return rest.NewMultiplexer(context.Background(), cfg)
}

// EOF
<|endoftext|>


<commit_before>package main

import (
	"github.com/hectane/go-asyncserver"

	"io"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

// Server acts as an HTTP proxy, returning entries from the cache whenever
// possible. Recognized archive mirrors are silently rewritten to avoid
// needless duplication.
type Server struct {
	server *server.AsyncServer
	cache  *Cache
}

func rewrite(rawurl string) string {
	u, err := url.Parse(rawurl)
	if err != nil {
		return rawurl
	}
	if strings.HasSuffix(u.Host, ".archive.ubuntu.com") {
		u.Host = "archive.ubuntu.com"
		rawurl = u.String()
	}
	return rawurl
}

func (s *Server) writeHeaders(w http.ResponseWriter, e *Entry) {
	if e.ContentType != "" {
		w.Header().Set("Content-Type", e.ContentType)
	} else {
		w.Header().Set("Content-Type", "application/octet-stream")
	}
	l, err := strconv.ParseInt(e.ContentLength, 10, 64)
	if err == nil && l >= 0 {
		w.Header().Set("Content-Length", e.ContentLength)
	}
	if e.LastModified != "" {
		w.Header().Set("Last-Modified", e.LastModified)
	}
	w.WriteHeader(http.StatusOK)
}

// TODO: support for HEAD requests
// TODO: find a reasonable way for getting errors from eChan

// ServeHTTP processes an incoming request to the proxy. GET requests are
// served with the storage backend and every other request is (out of
// necessity) rejected since it can't be cached.
func (s *Server) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if req.Method == "GET" {
		r, eChan, err := s.cache.GetReader(rewrite(req.RequestURI))
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		defer r.Close()
		e := <-eChan
		if e == nil {
			http.Error(w, "header retrieval error", http.StatusInternalServerError)
			return
		}
		s.writeHeaders(w, e)
		_, err = io.Copy(w, r)
		if err != nil {
			log.Println("[ERR]", err)
		}
	} else {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

// NewServer creates a new server.
func NewServer(addr, directory string) (*Server, error) {
	c, err := NewCache(directory)
	if err != nil {
		return nil, err
	}
	s := &Server{
		server: server.New(addr),
		cache:  c,
	}
	s.server.Handler = s
	return s, nil
}

// Start initializes the server and begins listening for requests.
func (s *Server) Start() error {
	return s.server.Start()
}

// Stop shuts down the server.
func (s *Server) Stop() {
	s.server.Stop()
	s.cache.Close()
}
<commit_msg>Added a bit more debug information to failed requests.<commit_after>package main

import (
	"github.com/hectane/go-asyncserver"

	"io"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

// Server acts as an HTTP proxy, returning entries from the cache whenever
// possible. Recognized archive mirrors are silently rewritten to avoid
// needless duplication.
type Server struct {
	server *server.AsyncServer
	cache  *Cache
}

func rewrite(rawurl string) string {
	u, err := url.Parse(rawurl)
	if err != nil {
		return rawurl
	}
	if strings.HasSuffix(u.Host, ".archive.ubuntu.com") {
		u.Host = "archive.ubuntu.com"
		rawurl = u.String()
	}
	return rawurl
}

func (s *Server) writeHeaders(w http.ResponseWriter, e *Entry) {
	if e.ContentType != "" {
		w.Header().Set("Content-Type", e.ContentType)
	} else {
		w.Header().Set("Content-Type", "application/octet-stream")
	}
	l, err := strconv.ParseInt(e.ContentLength, 10, 64)
	if err == nil && l >= 0 {
		w.Header().Set("Content-Length", e.ContentLength)
	}
	if e.LastModified != "" {
		w.Header().Set("Last-Modified", e.LastModified)
	}
	w.WriteHeader(http.StatusOK)
}

// TODO: support for HEAD requests
// TODO: find a reasonable way for getting errors from eChan

// ServeHTTP processes an incoming request to the proxy. GET requests are
// served with the storage backend and every other request is (out of
// necessity) rejected since it can't be cached.
func (s *Server) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if req.Method == "GET" {
		r, eChan, err := s.cache.GetReader(rewrite(req.RequestURI))
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			log.Println("[ERR]", err)
			return
		}
		defer r.Close()
		e := <-eChan
		if e == nil {
			http.Error(w, "header retrieval error", http.StatusInternalServerError)
			log.Println("[ERR] header retrieval")
			return
		}
		s.writeHeaders(w, e)
		_, err = io.Copy(w, r)
		if err != nil {
			log.Println("[ERR]", err)
		}
	} else {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

// NewServer creates a new server.
func NewServer(addr, directory string) (*Server, error) {
	c, err := NewCache(directory)
	if err != nil {
		return nil, err
	}
	s := &Server{
		server: server.New(addr),
		cache:  c,
	}
	s.server.Handler = s
	return s, nil
}

// Start initializes the server and begins listening for requests.
func (s *Server) Start() error {
	return s.server.Start()
}

// Stop shuts down the server.
func (s *Server) Stop() {
	s.server.Stop()
	s.cache.Close()
}
<|endoftext|>


<commit_before>package quic

import (
	"bytes"
	"crypto/tls"
	"net"
	"strings"
	"sync"

	"github.com/lucas-clemente/quic-go/crypto"
	"github.com/lucas-clemente/quic-go/handshake"
	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/qerr"
	"github.com/lucas-clemente/quic-go/utils"
)

// packetHandler handles packets
type packetHandler interface {
	handlePacket(addr interface{}, hdr *publicHeader, data []byte)
	run()
	Close(error) error
}

// A Server of QUIC
type Server struct {
	addr *net.UDPAddr
	conn *net.UDPConn

	signer crypto.Signer
	scfg   *handshake.ServerConfig

	sessions      map[protocol.ConnectionID]packetHandler
	sessionsMutex sync.RWMutex

	streamCallback StreamCallback

	newSession func(conn connection, v protocol.VersionNumber, connectionID protocol.ConnectionID, sCfg *handshake.ServerConfig, streamCallback StreamCallback, closeCallback closeCallback) (packetHandler, error)
}

// NewServer makes a new server
func NewServer(addr string, tlsConfig *tls.Config, cb StreamCallback) (*Server, error) {
	signer, err := crypto.NewProofSource(tlsConfig)
	if err != nil {
		return nil, err
	}

	kex, err := crypto.NewCurve25519KEX()
	if err != nil {
		return nil, err
	}
	scfg, err := handshake.NewServerConfig(kex, signer)
	if err != nil {
		return nil, err
	}

	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}

	return &Server{
		addr:           udpAddr,
		signer:         signer,
		scfg:           scfg,
		streamCallback: cb,
		sessions:       map[protocol.ConnectionID]packetHandler{},
		newSession:     newSession,
	}, nil
}

// ListenAndServe listens and serves a connection
func (s *Server) ListenAndServe() error {
	conn, err := net.ListenUDP("udp", s.addr)
	if err != nil {
		return err
	}
	s.conn = conn

	for {
		data := make([]byte, protocol.MaxPacketSize)
		n, remoteAddr, err := conn.ReadFromUDP(data)
		if err != nil {
			if strings.HasSuffix(err.Error(), "use of closed network connection") {
				return nil
			}
			return err
		}
		data = data[:n]
		if err := s.handlePacket(conn, remoteAddr, data); err != nil {
			utils.Errorf("error handling packet: %s", err.Error())
		}
	}
}

// Close the server
func (s *Server) Close() error {
	s.sessionsMutex.Lock()
	for _, session := range s.sessions {
		if session != nil {
			s.sessionsMutex.Unlock()
			_ = session.Close(nil)
			s.sessionsMutex.Lock()
		}
	}
	s.sessionsMutex.Unlock()
	if s.conn == nil {
		return nil
	}
	defer func() {
		s.conn = nil
	}()
	return s.conn.Close()
}

func (s *Server) handlePacket(conn *net.UDPConn, remoteAddr *net.UDPAddr, packet []byte) error {
	if protocol.ByteCount(len(packet)) > protocol.MaxPacketSize {
		return qerr.PacketTooLarge
	}

	r := bytes.NewReader(packet)

	hdr, err := parsePublicHeader(r)
	if err != nil {
		return qerr.Error(qerr.InvalidPacketHeader, err.Error())
	}
	hdr.Raw = packet[:len(packet)-r.Len()]

	// Send Version Negotiation Packet if the client is speaking a different protocol version
	if hdr.VersionFlag && !protocol.IsSupportedVersion(hdr.VersionNumber) {
		utils.Infof("Client offered version %d, sending VersionNegotiationPacket", hdr.VersionNumber)
		_, err = conn.WriteToUDP(composeVersionNegotiation(hdr.ConnectionID), remoteAddr)
		if err != nil {
			return err
		}
		return nil
	}

	s.sessionsMutex.RLock()
	session, ok := s.sessions[hdr.ConnectionID]
	s.sessionsMutex.RUnlock()

	if !ok {
		utils.Infof("Serving new connection: %x, version %d from %v", hdr.ConnectionID, hdr.VersionNumber, remoteAddr)
		session, err = s.newSession(
			&udpConn{conn: conn, currentAddr: remoteAddr},
			hdr.VersionNumber,
			hdr.ConnectionID,
			s.scfg,
			s.streamCallback,
			s.closeCallback,
		)
		if err != nil {
			return err
		}
		go session.run()
		s.sessionsMutex.Lock()
		s.sessions[hdr.ConnectionID] = session
		s.sessionsMutex.Unlock()
	}
	if session == nil {
		// Late packet for closed session
		return nil
	}
	session.handlePacket(remoteAddr, hdr, packet[len(packet)-r.Len():])
	return nil
}

func (s *Server) closeCallback(id protocol.ConnectionID) {
	s.sessionsMutex.Lock()
	s.sessions[id] = nil
	s.sessionsMutex.Unlock()
}

func composeVersionNegotiation(connectionID protocol.ConnectionID) []byte {
	fullReply := &bytes.Buffer{}
	responsePublicHeader := publicHeader{
		ConnectionID: connectionID,
		PacketNumber: 1,
		VersionFlag:  true,
	}
	// TODO: Update version number
	err := responsePublicHeader.WritePublicHeader(fullReply, protocol.VersionNumber(32))
	if err != nil {
		utils.Errorf("error composing version negotiation packet: %s", err.Error())
	}
	fullReply.Write(protocol.SupportedVersionsAsTags)
	return fullReply.Bytes()
}
<commit_msg>fix a race condition in quic.Server<commit_after>package quic

import (
	"bytes"
	"crypto/tls"
	"net"
	"strings"
	"sync"

	"github.com/lucas-clemente/quic-go/crypto"
	"github.com/lucas-clemente/quic-go/handshake"
	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/qerr"
	"github.com/lucas-clemente/quic-go/utils"
)

// packetHandler handles packets
type packetHandler interface {
	handlePacket(addr interface{}, hdr *publicHeader, data []byte)
	run()
	Close(error) error
}

// A Server of QUIC
type Server struct {
	addr *net.UDPAddr

	conn      *net.UDPConn
	connMutex sync.Mutex

	signer crypto.Signer
	scfg   *handshake.ServerConfig

	sessions      map[protocol.ConnectionID]packetHandler
	sessionsMutex sync.RWMutex

	streamCallback StreamCallback

	newSession func(conn connection, v protocol.VersionNumber, connectionID protocol.ConnectionID, sCfg *handshake.ServerConfig, streamCallback StreamCallback, closeCallback closeCallback) (packetHandler, error)
}

// NewServer makes a new server
func NewServer(addr string, tlsConfig *tls.Config, cb StreamCallback) (*Server, error) {
	signer, err := crypto.NewProofSource(tlsConfig)
	if err != nil {
		return nil, err
	}

	kex, err := crypto.NewCurve25519KEX()
	if err != nil {
		return nil, err
	}
	scfg, err := handshake.NewServerConfig(kex, signer)
	if err != nil {
		return nil, err
	}

	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}

	return &Server{
		addr:           udpAddr,
		signer:         signer,
		scfg:           scfg,
		streamCallback: cb,
		sessions:       map[protocol.ConnectionID]packetHandler{},
		newSession:     newSession,
	}, nil
}

// ListenAndServe listens and serves a connection
func (s *Server) ListenAndServe() error {
	conn, err := net.ListenUDP("udp", s.addr)
	if err != nil {
		return err
	}
	s.connMutex.Lock()
	s.conn = conn
	s.connMutex.Unlock()

	for {
		data := make([]byte, protocol.MaxPacketSize)
		n, remoteAddr, err := conn.ReadFromUDP(data)
		if err != nil {
			if strings.HasSuffix(err.Error(), "use of closed network connection") {
				return nil
			}
			return err
		}
		data = data[:n]
		if err := s.handlePacket(conn, remoteAddr, data); err != nil {
			utils.Errorf("error handling packet: %s", err.Error())
		}
	}
}

// Close the server
func (s *Server) Close() error {
	s.sessionsMutex.Lock()
	for _, session := range s.sessions {
		if session != nil {
			s.sessionsMutex.Unlock()
			_ = session.Close(nil)
			s.sessionsMutex.Lock()
		}
	}
	s.sessionsMutex.Unlock()

	s.connMutex.Lock()
	defer s.connMutex.Unlock()

	if s.conn == nil {
		return nil
	}
	defer func() {
		s.conn = nil
	}()
	return s.conn.Close()
}

func (s *Server) handlePacket(conn *net.UDPConn, remoteAddr *net.UDPAddr, packet []byte) error {
	if protocol.ByteCount(len(packet)) > protocol.MaxPacketSize {
		return qerr.PacketTooLarge
	}

	r := bytes.NewReader(packet)

	hdr, err := parsePublicHeader(r)
	if err != nil {
		return qerr.Error(qerr.InvalidPacketHeader, err.Error())
	}
	hdr.Raw = packet[:len(packet)-r.Len()]

	// Send Version Negotiation Packet if the client is speaking a different protocol version
	if hdr.VersionFlag && !protocol.IsSupportedVersion(hdr.VersionNumber) {
		utils.Infof("Client offered version %d, sending VersionNegotiationPacket", hdr.VersionNumber)
		_, err = conn.WriteToUDP(composeVersionNegotiation(hdr.ConnectionID), remoteAddr)
		if err != nil {
			return err
		}
		return nil
	}

	s.sessionsMutex.RLock()
	session, ok := s.sessions[hdr.ConnectionID]
	s.sessionsMutex.RUnlock()

	if !ok {
		utils.Infof("Serving new connection: %x, version %d from %v", hdr.ConnectionID, hdr.VersionNumber, remoteAddr)
		session, err = s.newSession(
			&udpConn{conn: conn, currentAddr: remoteAddr},
			hdr.VersionNumber,
			hdr.ConnectionID,
			s.scfg,
			s.streamCallback,
			s.closeCallback,
		)
		if err != nil {
			return err
		}
		go session.run()
		s.sessionsMutex.Lock()
		s.sessions[hdr.ConnectionID] = session
		s.sessionsMutex.Unlock()
	}
	if session == nil {
		// Late packet for closed session
		return nil
	}
	session.handlePacket(remoteAddr, hdr, packet[len(packet)-r.Len():])
	return nil
}

func (s *Server) closeCallback(id protocol.ConnectionID) {
	s.sessionsMutex.Lock()
	s.sessions[id] = nil
	s.sessionsMutex.Unlock()
}

func composeVersionNegotiation(connectionID protocol.ConnectionID) []byte {
	fullReply := &bytes.Buffer{}
	responsePublicHeader := publicHeader{
		ConnectionID: connectionID,
		PacketNumber: 1,
		VersionFlag:  true,
	}
	// TODO: Update version number
	err := responsePublicHeader.WritePublicHeader(fullReply, protocol.VersionNumber(32))
	if err != nil {
		utils.Errorf("error composing version negotiation packet: %s", err.Error())
	}
	fullReply.Write(protocol.SupportedVersionsAsTags)
	return fullReply.Bytes()
}
<|endoftext|>


<commit_before>package delayd

import (
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/streadway/amqp"
)

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

// Server is the delayd server. It handles the server lifecycle (startup, clean shutdown)
type Server struct {
	sender     Sender
	receiver   Receiver
	raft       *Raft
	timer      *Timer
	shutdownCh chan bool
	leader     bool
	mu         sync.Mutex
}

// NewServer initialize Server instance.
func NewServer(c Config) (*Server, error) {
	if len(c.LogDir) != 0 {
		logFile := filepath.Join(c.LogDir, "delayd.log")
		logOutput, err := os.OpenFile(logFile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
		if err != nil {
			return nil, err
		}
		log.SetOutput(logOutput)
	}

	receiver, err := NewAMQPReceiver(c.AMQP, routingKey)
	if err != nil {
		return nil, err
	}

	sender, err := NewAMQPSender(c.AMQP.URL)
	if err != nil {
		return nil, err
	}

	raft, err := NewRaft(c.Raft, c.DataDir, c.LogDir)
	if err != nil {
		return nil, err
	}

	return &Server{
		sender:     sender,
		receiver:   receiver,
		raft:       raft,
		shutdownCh: make(chan bool),
	}, nil
}

// Run starts server and begins its main loop.
func (s *Server) Run() {
	Info("server: starting delayd")

	s.timer = NewTimer(s.timerSend)

	go s.observeLeaderChanges()
	go s.observeNextTime()

	for {
		msg, ok := <-s.receiver.MessageCh()
		entry := msg.Entry
		// XXX cleanup needed here before exit
		if !ok {
			continue
		}

		Debug("server: got new request entry:", entry)
		b, err := entry.ToBytes()
		if err != nil {
			Error("server: error encoding entry: ", err)
			continue
		}

		if err := s.raft.Add(b, raftMaxTime); err != nil {
			Error("server: failed to add: ", err)
			msg.Nack()
			continue
		}

		msg.Ack()
	}
}

// Stop shuts down the Server cleanly. Order of the Close calls is important.
func (s *Server) Stop() {
	Info("server: shutting down gracefully.")

	// stop triggering new changes to the FSM
	s.receiver.Close()
	s.timer.Stop()

	close(s.shutdownCh)

	// with no FSM stimulus, we don't need to send.
	s.sender.Close()

	s.raft.Close()

	Info("server: terminated.")
}

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

// Listen for changes to raft for when we transition to and from the leader state
// and react accordingly.
func (s *Server) observeLeaderChanges() {
	for isLeader := range s.raft.LeaderCh() {
		s.mu.Lock()
		s.leader = isLeader
		s.mu.Unlock()

		if isLeader {
			Debug("server: became raft leader")
			s.resetTimer()
			err := s.receiver.Start()
			if err != nil {
				Fatal("server: error while starting receiver:", err)
			}
		} else {
			Debug("server: lost raft leadership")
			s.timer.Pause()
			err := s.receiver.Pause()
			if err != nil {
				Fatal("server: error while starting receiver:", err)
			}
		}
	}
}

// Listen to storage for next time changes on entry commits
func (s *Server) observeNextTime() {
	for {
		select {
		case <-s.shutdownCh:
			return
		case t := <-s.raft.fsm.store.C:
			s.mu.Lock()
			if s.leader {
				Debug("server: got time: ", t)
				s.timer.Reset(t, false)
			}
			s.mu.Unlock()
		}
	}
}

func (s *Server) timerSend(t time.Time) {
	// FIXME: In the case of error, we don't remove a log from raft log.
	uuids, entries, err := s.raft.fsm.store.Get(t)
	if err != nil {
		Fatal("server: could not read entries from db:", err)
	}

	Infof("server: sending %d entries", len(entries))
	if len(entries) < 1 {
		// FIXME: if no entries found, timer is missing newest entries.........
		_, entries, _ := s.raft.fsm.store.GetAll()
		Warnf("server: timer is missing newest entries. total: %d", len(entries))
	}

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

		if err := s.sender.Send(e); err != nil {
			if err, ok := err.(*amqp.Error); ok && err.Code == 504 {
				Warnf("server: channel/connection not set up for exchange `%s`, message will be deleted", e.Target, err)
				break
			} else {
				Fatal("server: could not send entry: ", err)
			}
		}

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

	// ensure everyone is up to date
	Debug("server: syncing raft after send.")
	err = s.raft.SyncAll()
	if err != nil {
		Warn("server: lost raft leadership during sync after send.")
		return
	}

	Debug("server: synced raft after send.")
}
<commit_msg>server: cleanup logs and do not break if err.Code == 504<commit_after>package delayd

import (
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/streadway/amqp"
)

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

// Server is the delayd server. It handles the server lifecycle (startup, clean shutdown)
type Server struct {
	sender     Sender
	receiver   Receiver
	raft       *Raft
	timer      *Timer
	shutdownCh chan bool
	leader     bool
	mu         sync.Mutex
}

// NewServer initialize Server instance.
func NewServer(c Config) (*Server, error) {
	if len(c.LogDir) != 0 {
		logFile := filepath.Join(c.LogDir, "delayd.log")
		logOutput, err := os.OpenFile(logFile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
		if err != nil {
			return nil, err
		}
		log.SetOutput(logOutput)
	}

	receiver, err := NewAMQPReceiver(c.AMQP, routingKey)
	if err != nil {
		return nil, err
	}

	sender, err := NewAMQPSender(c.AMQP.URL)
	if err != nil {
		return nil, err
	}

	raft, err := NewRaft(c.Raft, c.DataDir, c.LogDir)
	if err != nil {
		return nil, err
	}

	return &Server{
		sender:     sender,
		receiver:   receiver,
		raft:       raft,
		shutdownCh: make(chan bool),
	}, nil
}

// Run starts server and begins its main loop.
func (s *Server) Run() {
	Info("server: starting delayd")

	s.timer = NewTimer(s.timerSend)

	go s.observeLeaderChanges()
	go s.observeNextTime()

	for {
		msg, ok := <-s.receiver.MessageCh()
		entry := msg.Entry
		// XXX cleanup needed here before exit
		if !ok {
			continue
		}

		Debug("server: got new request entry:", entry)
		b, err := entry.ToBytes()
		if err != nil {
			Error("server: error encoding entry: ", err)
			continue
		}

		if err := s.raft.Add(b, raftMaxTime); err != nil {
			Error("server: failed to add: ", err)
			msg.Nack()
			continue
		}

		msg.Ack()
	}
}

// Stop shuts down the Server cleanly. Order of the Close calls is important.
func (s *Server) Stop() {
	Info("server: shutting down gracefully.")

	// stop triggering new changes to the FSM
	s.receiver.Close()
	s.timer.Stop()

	close(s.shutdownCh)

	// with no FSM stimulus, we don't need to send.
	s.sender.Close()

	s.raft.Close()

	Info("server: terminated.")
}

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

// Listen for changes to raft for when we transition to and from the leader state
// and react accordingly.
func (s *Server) observeLeaderChanges() {
	for isLeader := range s.raft.LeaderCh() {
		s.mu.Lock()
		s.leader = isLeader
		s.mu.Unlock()

		if isLeader {
			Debug("server: became raft leader")
			s.resetTimer()
			err := s.receiver.Start()
			if err != nil {
				Fatal("server: error while starting receiver:", err)
			}
		} else {
			Debug("server: lost raft leadership")
			s.timer.Pause()
			err := s.receiver.Pause()
			if err != nil {
				Fatal("server: error while starting receiver:", err)
			}
		}
	}
}

// Listen to storage for next time changes on entry commits
func (s *Server) observeNextTime() {
	for {
		select {
		case <-s.shutdownCh:
			return
		case t := <-s.raft.fsm.store.C:
			s.mu.Lock()
			if s.leader {
				Debug("server: got time: ", t)
				s.timer.Reset(t, false)
			}
			s.mu.Unlock()
		}
	}
}

func (s *Server) timerSend(t time.Time) {
	// FIXME: In the case of error, we don't remove a log from raft log.
	uuids, entries, err := s.raft.fsm.store.Get(t)
	if err != nil {
		Fatal("server: could not read entries from db:", err)
	}

	if len(entries) > 1 {
		Infof("server: sending %d entries", len(entries))
	}

	for i, e := range entries {
		if err := s.sender.Send(e); err != nil {
			if err, ok := err.(*amqp.Error); ok && err.Code == 504 {
				// error 504 code means that the exchange we were trying
				// to send on didnt exist.  In the case of delayd this usually
				// means that a consumer didn't set up the exchange they wish
				// to be notified on. We do not attempt to make this for them,
				// as we don't know what exchange options they would want, we
				// simply drop this message, other errors are fatal
				Warnf("server: channel/connection not set up for exchange `%s`, message will be deleted", e.Target, err)
			} else {
				Fatal("server: could not send entry:", err)
			}
		}

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

	// ensure everyone is up to date
	Debug("server: syncing raft after send.")
	err = s.raft.SyncAll()
	if err != nil {
		Warn("server: lost raft leadership during sync after send.")
		return
	}

	Debug("server: synced raft after send.")
}
<|endoftext|>


<commit_before>package main

import (
  "github.com/go-martini/martini"
  "github.com/martini-contrib/binding"
  "github.com/martini-contrib/render"
  "labix.org/v2/mgo"
)

type Quote struct {
  Content   string  `form:"content"`
}

func DB() martini.Handler {
  session, err := mgo.Dial("mongodb://localhost")
  if err != nil {
    panic(err)
  }

  return func(c martini.Context) {
    s := session.Clone()
    c.Map(s.DB("quotes"))
  }
}

func GetAll(db *mgo.Database) []Quote {
  var quotes []Quote
  db.C("quotes").Find(nil).All(&quotes)
  return quotes
}

func main() {
  m := martini.Classic()

  // render html templates from templates directory
  m.Use(render.Renderer(render.Options{
    Layout: "layout",
  }))

  m.Use(DB())

  m.Get("/", func(r render.Render, db *mgo.Database) {
    data := map[string]interface{}{"quotes": GetAll(db)}
    r.HTML(200, "list", data)
  })

  m.Post("/", binding.Form(Quote{}), func(r render.Render, db *mgo.Database, quote Quote) {
    db.C("quotes").Insert(quote)
    data := map[string]interface{}{"quotes": GetAll(db)}
    r.HTML(200, "list", data)
  })

  m.Run()
}
<commit_msg>Added gzip<commit_after>package main

import (
  "github.com/go-martini/martini"
  "github.com/martini-contrib/binding"
  "github.com/martini-contrib/render"
  "github.com/martini-contrib/gzip"
  "labix.org/v2/mgo"
)

type Quote struct {
  Content   string  `form:"content"`
}

func DB() martini.Handler {
  session, err := mgo.Dial("mongodb://localhost")
  if err != nil {
    panic(err)
  }

  return func(c martini.Context) {
    s := session.Clone()
    c.Map(s.DB("quotes"))
  }
}

func GetAll(db *mgo.Database) []Quote {
  var quotes []Quote
  db.C("quotes").Find(nil).All(&quotes)
  return quotes
}

func main() {
  m := martini.Classic()

  m.Use(gzip.All())

  // render html templates from templates directory
  m.Use(render.Renderer(render.Options{
    Layout: "layout",
  }))

  m.Use(DB())

  m.Get("/", func(r render.Render, db *mgo.Database) {
    data := map[string]interface{}{"quotes": GetAll(db)}
    r.HTML(200, "list", data)
  })

  m.Post("/", binding.Form(Quote{}), func(r render.Render, db *mgo.Database, quote Quote) {
    db.C("quotes").Insert(quote)
    data := map[string]interface{}{"quotes": GetAll(db)}
    r.HTML(200, "list", data)
  })

  m.Run()
}
<|endoftext|>