<commit_before>package quic

import (
	"bytes"
	"crypto/tls"
	"errors"
	"net"
	"sync"
	"time"

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

// packetHandler handles packets
type packetHandler interface {
	Session
	handlePacket(*receivedPacket)
	GetVersion() protocol.VersionNumber
	run() error
	closeRemote(error)
}

// A Listener of QUIC
type server struct {
	tlsConf *tls.Config
	config  *Config

	conn net.PacketConn

	certChain crypto.CertChain
	scfg      *handshake.ServerConfig

	sessions                  map[protocol.ConnectionID]packetHandler
	sessionsMutex             sync.RWMutex
	deleteClosedSessionsAfter time.Duration

	serverError  error
	sessionQueue chan Session
	errorChan    chan struct{}

	newSession func(conn connection, v protocol.VersionNumber, connectionID protocol.ConnectionID, sCfg *handshake.ServerConfig, tlsConf *tls.Config, config *Config) (packetHandler, <-chan handshakeEvent, error)
}

var _ Listener = &server{}

// ListenAddr creates a QUIC server listening on a given address.
// The listener is not active until Serve() is called.
// The tls.Config must not be nil, the quic.Config may be nil.
func ListenAddr(addr string, tlsConf *tls.Config, config *Config) (Listener, error) {
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return nil, err
	}
	return Listen(conn, tlsConf, config)
}

// Listen listens for QUIC connections on a given net.PacketConn.
// The listener is not active until Serve() is called.
// The tls.Config must not be nil, the quic.Config may be nil.
func Listen(conn net.PacketConn, tlsConf *tls.Config, config *Config) (Listener, error) {
	certChain := crypto.NewCertChain(tlsConf)
	kex, err := crypto.NewCurve25519KEX()
	if err != nil {
		return nil, err
	}
	scfg, err := handshake.NewServerConfig(kex, certChain)
	if err != nil {
		return nil, err
	}

	s := &server{
		conn:                      conn,
		tlsConf:                   tlsConf,
		config:                    populateServerConfig(config),
		certChain:                 certChain,
		scfg:                      scfg,
		sessions:                  map[protocol.ConnectionID]packetHandler{},
		newSession:                newSession,
		deleteClosedSessionsAfter: protocol.ClosedSessionDeleteTimeout,
		sessionQueue:              make(chan Session, 5),
		errorChan:                 make(chan struct{}),
	}
	go s.serve()
	utils.Debugf("Listening for %s connections on %s", conn.LocalAddr().Network(), conn.LocalAddr().String())
	return s, nil
}

var defaultAcceptCookie = func(clientAddr net.Addr, cookie *Cookie) bool {
	if cookie == nil {
		return false
	}
	if time.Now().After(cookie.SentTime.Add(protocol.CookieExpiryTime)) {
		return false
	}
	var sourceAddr string
	if udpAddr, ok := clientAddr.(*net.UDPAddr); ok {
		sourceAddr = udpAddr.IP.String()
	} else {
		sourceAddr = clientAddr.String()
	}
	return sourceAddr == cookie.RemoteAddr
}

// populateServerConfig populates fields in the quic.Config with their default values, if none are set
// it may be called with nil
func populateServerConfig(config *Config) *Config {
	if config == nil {
		config = &Config{}
	}
	versions := config.Versions
	if len(versions) == 0 {
		versions = protocol.SupportedVersions
	}

	vsa := defaultAcceptCookie
	if config.AcceptCookie != nil {
		vsa = config.AcceptCookie
	}

	handshakeTimeout := protocol.DefaultHandshakeTimeout
	if config.HandshakeTimeout != 0 {
		handshakeTimeout = config.HandshakeTimeout
	}
	idleTimeout := protocol.DefaultIdleTimeout
	if config.IdleTimeout != 0 {
		idleTimeout = config.IdleTimeout
	}

	maxReceiveStreamFlowControlWindow := config.MaxReceiveStreamFlowControlWindow
	if maxReceiveStreamFlowControlWindow == 0 {
		maxReceiveStreamFlowControlWindow = protocol.DefaultMaxReceiveStreamFlowControlWindowServer
	}
	maxReceiveConnectionFlowControlWindow := config.MaxReceiveConnectionFlowControlWindow
	if maxReceiveConnectionFlowControlWindow == 0 {
		maxReceiveConnectionFlowControlWindow = protocol.DefaultMaxReceiveConnectionFlowControlWindowServer
	}

	return &Config{
		Versions:                              versions,
		HandshakeTimeout:                      handshakeTimeout,
		IdleTimeout:                           idleTimeout,
		AcceptCookie:                          vsa,
		KeepAlive:                             config.KeepAlive,
		MaxReceiveStreamFlowControlWindow:     maxReceiveStreamFlowControlWindow,
		MaxReceiveConnectionFlowControlWindow: maxReceiveConnectionFlowControlWindow,
	}
}

// serve listens on an existing PacketConn
func (s *server) serve() {
	for {
		data := getPacketBuffer()
		data = data[:protocol.MaxReceivePacketSize]
		// The packet size should not exceed protocol.MaxReceivePacketSize bytes
		// If it does, we only read a truncated packet, which will then end up undecryptable
		n, remoteAddr, err := s.conn.ReadFrom(data)
		if err != nil {
			s.serverError = err
			close(s.errorChan)
			_ = s.Close()
			return
		}
		data = data[:n]
		if err := s.handlePacket(s.conn, remoteAddr, data); err != nil {
			utils.Errorf("error handling packet: %s", err.Error())
		}
	}
}

// Accept returns newly openend sessions
func (s *server) Accept() (Session, error) {
	var sess Session
	select {
	case sess = <-s.sessionQueue:
		return sess, nil
	case <-s.errorChan:
		return nil, s.serverError
	}
}

// Close the server
func (s *server) Close() error {
	s.sessionsMutex.Lock()
	var wg sync.WaitGroup
	for _, session := range s.sessions {
		if session != nil {
			wg.Add(1)
			go func() {
				// session.Close() blocks until the CONNECTION_CLOSE has been sent and the run-loop has stopped
				_ = session.Close(nil)
				wg.Done()
			}()
		}
	}
	s.sessionsMutex.Unlock()
	wg.Wait()

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

// Addr returns the server's network address
func (s *server) Addr() net.Addr {
	return s.conn.LocalAddr()
}

func (s *server) handlePacket(pconn net.PacketConn, remoteAddr net.Addr, packet []byte) error {
	rcvTime := time.Now()

	r := bytes.NewReader(packet)
	connID, err := wire.PeekConnectionID(r)
	if err != nil {
		return qerr.Error(qerr.InvalidPacketHeader, err.Error())
	}

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

	if ok && session == nil {
		// Late packet for closed session
		return nil
	}

	version := protocol.VersionUnknown
	if ok {
		version = session.GetVersion()
	}

	hdr, err := wire.ParseHeader(r, protocol.PerspectiveClient, version)
	if err == wire.ErrPacketWithUnknownVersion {
		_, err = pconn.WriteTo(wire.WritePublicReset(connID, 0, 0), remoteAddr)
		return err
	}
	if err != nil {
		return qerr.Error(qerr.InvalidPacketHeader, err.Error())
	}
	hdr.Raw = packet[:len(packet)-r.Len()]

	// ignore all Public Reset packets
	if hdr.ResetFlag {
		if ok {
			var pr *wire.PublicReset
			pr, err = wire.ParsePublicReset(r)
			if err != nil {
				utils.Infof("Received a Public Reset for connection %x. An error occurred parsing the packet.")
			} else {
				utils.Infof("Received a Public Reset for connection %x, rejected packet number: 0x%x.", hdr.ConnectionID, pr.RejectedPacketNumber)
			}
		} else {
			utils.Infof("Received Public Reset for unknown connection %x.", hdr.ConnectionID)
		}
		return nil
	}

	// a session is only created once the client sent a supported version
	// if we receive a packet for a connection that already has session, it's probably an old packet that was sent by the client before the version was negotiated
	// it is safe to drop it
	if ok && hdr.VersionFlag && !protocol.IsSupportedVersion(s.config.Versions, hdr.Version) {
		return nil
	}

	// send a Version Negotiation Packet if the client is speaking a different protocol version
	// since the client send a Public Header (only gQUIC has a Version Flag), we need to send a gQUIC Version Negotiation Packet
	if hdr.VersionFlag && !protocol.IsSupportedVersion(s.config.Versions, hdr.Version) {
		// drop packets that are too small to be valid first packets
		if len(packet) < protocol.ClientHelloMinimumSize+len(hdr.Raw) {
			return errors.New("dropping small packet with unknown version")
		}
		utils.Infof("Client offered version %s, sending VersionNegotiationPacket", hdr.Version)
		if _, err := pconn.WriteTo(wire.ComposeGQUICVersionNegotiation(hdr.ConnectionID, s.config.Versions), remoteAddr); err != nil {
			return err
		}
	}
	// send an IETF draft style Version Negotiation Packet, if the client sent an unsupported version with an IETF draft style header
	if hdr.Type == protocol.PacketTypeInitial && !protocol.IsSupportedVersion(s.config.Versions, hdr.Version) {
		_, err := pconn.WriteTo(wire.ComposeVersionNegotiation(hdr.ConnectionID, hdr.PacketNumber, s.config.Versions), remoteAddr)
		return err
	}

	if !ok {
		version := hdr.Version
		if !protocol.IsSupportedVersion(s.config.Versions, version) {
			return errors.New("Server BUG: negotiated version not supported")
		}

		utils.Infof("Serving new connection: %x, version %s from %v", hdr.ConnectionID, version, remoteAddr)
		var handshakeChan <-chan handshakeEvent
		session, handshakeChan, err = s.newSession(
			&conn{pconn: pconn, currentAddr: remoteAddr},
			version,
			hdr.ConnectionID,
			s.scfg,
			s.tlsConf,
			s.config,
		)
		if err != nil {
			return err
		}
		s.sessionsMutex.Lock()
		s.sessions[connID] = session
		s.sessionsMutex.Unlock()

		go func() {
			// session.run() returns as soon as the session is closed
			_ = session.run()
			s.removeConnection(connID)
		}()

		go func() {
			for {
				ev := <-handshakeChan
				if ev.err != nil {
					return
				}
				if ev.encLevel == protocol.EncryptionForwardSecure {
					break
				}
			}
			s.sessionQueue <- session
		}()
	}
	session.handlePacket(&receivedPacket{
		remoteAddr: remoteAddr,
		header:     hdr,
		data:       packet[len(packet)-r.Len():],
		rcvTime:    rcvTime,
	})
	return nil
}

func (s *server) removeConnection(id protocol.ConnectionID) {
	s.sessionsMutex.Lock()
	s.sessions[id] = nil
	s.sessionsMutex.Unlock()

	time.AfterFunc(s.deleteClosedSessionsAfter, func() {
		s.sessionsMutex.Lock()
		delete(s.sessions, id)
		s.sessionsMutex.Unlock()
	})
}
<commit_msg>fix race condition when closing the server<commit_after>package quic

import (
	"bytes"
	"crypto/tls"
	"errors"
	"net"
	"sync"
	"time"

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

// packetHandler handles packets
type packetHandler interface {
	Session
	handlePacket(*receivedPacket)
	GetVersion() protocol.VersionNumber
	run() error
	closeRemote(error)
}

// A Listener of QUIC
type server struct {
	tlsConf *tls.Config
	config  *Config

	conn net.PacketConn

	certChain crypto.CertChain
	scfg      *handshake.ServerConfig

	sessions                  map[protocol.ConnectionID]packetHandler
	sessionsMutex             sync.RWMutex
	deleteClosedSessionsAfter time.Duration

	serverError  error
	sessionQueue chan Session
	errorChan    chan struct{}

	newSession func(conn connection, v protocol.VersionNumber, connectionID protocol.ConnectionID, sCfg *handshake.ServerConfig, tlsConf *tls.Config, config *Config) (packetHandler, <-chan handshakeEvent, error)
}

var _ Listener = &server{}

// ListenAddr creates a QUIC server listening on a given address.
// The listener is not active until Serve() is called.
// The tls.Config must not be nil, the quic.Config may be nil.
func ListenAddr(addr string, tlsConf *tls.Config, config *Config) (Listener, error) {
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return nil, err
	}
	return Listen(conn, tlsConf, config)
}

// Listen listens for QUIC connections on a given net.PacketConn.
// The listener is not active until Serve() is called.
// The tls.Config must not be nil, the quic.Config may be nil.
func Listen(conn net.PacketConn, tlsConf *tls.Config, config *Config) (Listener, error) {
	certChain := crypto.NewCertChain(tlsConf)
	kex, err := crypto.NewCurve25519KEX()
	if err != nil {
		return nil, err
	}
	scfg, err := handshake.NewServerConfig(kex, certChain)
	if err != nil {
		return nil, err
	}

	s := &server{
		conn:                      conn,
		tlsConf:                   tlsConf,
		config:                    populateServerConfig(config),
		certChain:                 certChain,
		scfg:                      scfg,
		sessions:                  map[protocol.ConnectionID]packetHandler{},
		newSession:                newSession,
		deleteClosedSessionsAfter: protocol.ClosedSessionDeleteTimeout,
		sessionQueue:              make(chan Session, 5),
		errorChan:                 make(chan struct{}),
	}
	go s.serve()
	utils.Debugf("Listening for %s connections on %s", conn.LocalAddr().Network(), conn.LocalAddr().String())
	return s, nil
}

var defaultAcceptCookie = func(clientAddr net.Addr, cookie *Cookie) bool {
	if cookie == nil {
		return false
	}
	if time.Now().After(cookie.SentTime.Add(protocol.CookieExpiryTime)) {
		return false
	}
	var sourceAddr string
	if udpAddr, ok := clientAddr.(*net.UDPAddr); ok {
		sourceAddr = udpAddr.IP.String()
	} else {
		sourceAddr = clientAddr.String()
	}
	return sourceAddr == cookie.RemoteAddr
}

// populateServerConfig populates fields in the quic.Config with their default values, if none are set
// it may be called with nil
func populateServerConfig(config *Config) *Config {
	if config == nil {
		config = &Config{}
	}
	versions := config.Versions
	if len(versions) == 0 {
		versions = protocol.SupportedVersions
	}

	vsa := defaultAcceptCookie
	if config.AcceptCookie != nil {
		vsa = config.AcceptCookie
	}

	handshakeTimeout := protocol.DefaultHandshakeTimeout
	if config.HandshakeTimeout != 0 {
		handshakeTimeout = config.HandshakeTimeout
	}
	idleTimeout := protocol.DefaultIdleTimeout
	if config.IdleTimeout != 0 {
		idleTimeout = config.IdleTimeout
	}

	maxReceiveStreamFlowControlWindow := config.MaxReceiveStreamFlowControlWindow
	if maxReceiveStreamFlowControlWindow == 0 {
		maxReceiveStreamFlowControlWindow = protocol.DefaultMaxReceiveStreamFlowControlWindowServer
	}
	maxReceiveConnectionFlowControlWindow := config.MaxReceiveConnectionFlowControlWindow
	if maxReceiveConnectionFlowControlWindow == 0 {
		maxReceiveConnectionFlowControlWindow = protocol.DefaultMaxReceiveConnectionFlowControlWindowServer
	}

	return &Config{
		Versions:                              versions,
		HandshakeTimeout:                      handshakeTimeout,
		IdleTimeout:                           idleTimeout,
		AcceptCookie:                          vsa,
		KeepAlive:                             config.KeepAlive,
		MaxReceiveStreamFlowControlWindow:     maxReceiveStreamFlowControlWindow,
		MaxReceiveConnectionFlowControlWindow: maxReceiveConnectionFlowControlWindow,
	}
}

// serve listens on an existing PacketConn
func (s *server) serve() {
	for {
		data := getPacketBuffer()
		data = data[:protocol.MaxReceivePacketSize]
		// The packet size should not exceed protocol.MaxReceivePacketSize bytes
		// If it does, we only read a truncated packet, which will then end up undecryptable
		n, remoteAddr, err := s.conn.ReadFrom(data)
		if err != nil {
			s.serverError = err
			close(s.errorChan)
			_ = s.Close()
			return
		}
		data = data[:n]
		if err := s.handlePacket(s.conn, remoteAddr, data); err != nil {
			utils.Errorf("error handling packet: %s", err.Error())
		}
	}
}

// Accept returns newly openend sessions
func (s *server) Accept() (Session, error) {
	var sess Session
	select {
	case sess = <-s.sessionQueue:
		return sess, nil
	case <-s.errorChan:
		return nil, s.serverError
	}
}

// Close the server
func (s *server) Close() error {
	s.sessionsMutex.Lock()
	var wg sync.WaitGroup
	for _, session := range s.sessions {
		if session != nil {
			wg.Add(1)
			go func(sess packetHandler) {
				// session.Close() blocks until the CONNECTION_CLOSE has been sent and the run-loop has stopped
				_ = sess.Close(nil)
				wg.Done()
			}(session)
		}
	}
	s.sessionsMutex.Unlock()
	wg.Wait()

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

// Addr returns the server's network address
func (s *server) Addr() net.Addr {
	return s.conn.LocalAddr()
}

func (s *server) handlePacket(pconn net.PacketConn, remoteAddr net.Addr, packet []byte) error {
	rcvTime := time.Now()

	r := bytes.NewReader(packet)
	connID, err := wire.PeekConnectionID(r)
	if err != nil {
		return qerr.Error(qerr.InvalidPacketHeader, err.Error())
	}

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

	if ok && session == nil {
		// Late packet for closed session
		return nil
	}

	version := protocol.VersionUnknown
	if ok {
		version = session.GetVersion()
	}

	hdr, err := wire.ParseHeader(r, protocol.PerspectiveClient, version)
	if err == wire.ErrPacketWithUnknownVersion {
		_, err = pconn.WriteTo(wire.WritePublicReset(connID, 0, 0), remoteAddr)
		return err
	}
	if err != nil {
		return qerr.Error(qerr.InvalidPacketHeader, err.Error())
	}
	hdr.Raw = packet[:len(packet)-r.Len()]

	// ignore all Public Reset packets
	if hdr.ResetFlag {
		if ok {
			var pr *wire.PublicReset
			pr, err = wire.ParsePublicReset(r)
			if err != nil {
				utils.Infof("Received a Public Reset for connection %x. An error occurred parsing the packet.")
			} else {
				utils.Infof("Received a Public Reset for connection %x, rejected packet number: 0x%x.", hdr.ConnectionID, pr.RejectedPacketNumber)
			}
		} else {
			utils.Infof("Received Public Reset for unknown connection %x.", hdr.ConnectionID)
		}
		return nil
	}

	// a session is only created once the client sent a supported version
	// if we receive a packet for a connection that already has session, it's probably an old packet that was sent by the client before the version was negotiated
	// it is safe to drop it
	if ok && hdr.VersionFlag && !protocol.IsSupportedVersion(s.config.Versions, hdr.Version) {
		return nil
	}

	// send a Version Negotiation Packet if the client is speaking a different protocol version
	// since the client send a Public Header (only gQUIC has a Version Flag), we need to send a gQUIC Version Negotiation Packet
	if hdr.VersionFlag && !protocol.IsSupportedVersion(s.config.Versions, hdr.Version) {
		// drop packets that are too small to be valid first packets
		if len(packet) < protocol.ClientHelloMinimumSize+len(hdr.Raw) {
			return errors.New("dropping small packet with unknown version")
		}
		utils.Infof("Client offered version %s, sending VersionNegotiationPacket", hdr.Version)
		if _, err := pconn.WriteTo(wire.ComposeGQUICVersionNegotiation(hdr.ConnectionID, s.config.Versions), remoteAddr); err != nil {
			return err
		}
	}
	// send an IETF draft style Version Negotiation Packet, if the client sent an unsupported version with an IETF draft style header
	if hdr.Type == protocol.PacketTypeInitial && !protocol.IsSupportedVersion(s.config.Versions, hdr.Version) {
		_, err := pconn.WriteTo(wire.ComposeVersionNegotiation(hdr.ConnectionID, hdr.PacketNumber, s.config.Versions), remoteAddr)
		return err
	}

	if !ok {
		version := hdr.Version
		if !protocol.IsSupportedVersion(s.config.Versions, version) {
			return errors.New("Server BUG: negotiated version not supported")
		}

		utils.Infof("Serving new connection: %x, version %s from %v", hdr.ConnectionID, version, remoteAddr)
		var handshakeChan <-chan handshakeEvent
		session, handshakeChan, err = s.newSession(
			&conn{pconn: pconn, currentAddr: remoteAddr},
			version,
			hdr.ConnectionID,
			s.scfg,
			s.tlsConf,
			s.config,
		)
		if err != nil {
			return err
		}
		s.sessionsMutex.Lock()
		s.sessions[connID] = session
		s.sessionsMutex.Unlock()

		go func() {
			// session.run() returns as soon as the session is closed
			_ = session.run()
			s.removeConnection(connID)
		}()

		go func() {
			for {
				ev := <-handshakeChan
				if ev.err != nil {
					return
				}
				if ev.encLevel == protocol.EncryptionForwardSecure {
					break
				}
			}
			s.sessionQueue <- session
		}()
	}
	session.handlePacket(&receivedPacket{
		remoteAddr: remoteAddr,
		header:     hdr,
		data:       packet[len(packet)-r.Len():],
		rcvTime:    rcvTime,
	})
	return nil
}

func (s *server) removeConnection(id protocol.ConnectionID) {
	s.sessionsMutex.Lock()
	s.sessions[id] = nil
	s.sessionsMutex.Unlock()

	time.AfterFunc(s.deleteClosedSessionsAfter, func() {
		s.sessionsMutex.Lock()
		delete(s.sessions, id)
		s.sessionsMutex.Unlock()
	})
}
<|endoftext|>


<commit_before>/*
 *  Glue - Robust Go and Javascript Socket Library
 *  Copyright (C) 2015  Roland Singer <roland.singer[at]desertbit.com>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

// Package glue - Robust Go and Javascript Socket Library.
// This library is thread-safe.
package glue

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

	"github.com/desertbit/glue/backend"
)

//####################//
//### Public Types ###//
//####################//

// OnNewSocketFunc is an event function.
type OnNewSocketFunc func(s *Socket)

//###################//
//### Server Type ###//
//###################//

// A Server represents a glue server which handles incoming socket connections.
type Server struct {
	bs      *backend.Server
	options *Options

	block       bool
	onNewSocket OnNewSocketFunc

	sockets      map[string]*Socket // A map holding all active current sockets.
	socketsMutex sync.Mutex
}

// NewServer creates a new glue server instance.
// One variadic arguments specifies the server options.
func NewServer(o ...Options) *Server {
	// Get or create the options.
	var options *Options
	if len(o) > 0 {
		options = &o[0]
	} else {
		options = &Options{}
	}

	// Set the default option values for unset values.
	options.SetDefaults()

	// Create a new backend server.
	bs := backend.NewServer(len(options.HTTPHandleURL))

	// Create a new server value.
	s := &Server{
		bs:          bs,
		options:     options,
		onNewSocket: func(*Socket) {}, // Initialize with dummy function to remove nil check.
		sockets:     make(map[string]*Socket),
	}

	// Set the backend server event function.
	bs.OnNewSocketConnection(s.handleOnNewSocketConnection)

	return s
}

// Block new incomming connections.
func (s *Server) Block(b bool) {
	s.block = b
}

// OnNewSocket sets the event function which is
// triggered if a new socket connection was made.
// The event function must not block! As soon as the event function
// returns, the socket is added to the active sockets map.
func (s *Server) OnNewSocket(f OnNewSocketFunc) {
	s.onNewSocket = f
}

// Sockets returns a list of all current connected sockets.
// Hint: Sockets are added to the active sockets list before the OnNewSocket
// event function is called.
// Use the IsInitialized flag to determind if a socket is not ready yet...
func (s *Server) Sockets() []*Socket {
	// Lock the mutex.
	s.socketsMutex.Lock()
	defer s.socketsMutex.Unlock()

	// Create the slice.
	list := make([]*Socket, len(s.sockets))

	// Add all sockets from the map.
	i := 0
	for _, s := range s.sockets {
		list[i] = s
		i++
	}

	return list
}

// Release this package. This will block all new incomming socket connections
// and close all current connected sockets.
func (s *Server) Release() {
	// Block all new incomming socket connections.
	s.Block(true)

	// Wait for a little moment, so new incomming sockets are added
	// to the sockets active list.
	time.Sleep(200 * time.Millisecond)

	// Close all current connected sockets.
	sockets := s.Sockets()
	for _, s := range sockets {
		s.Close()
	}
}

// Run starts the server and listens for incoming socket connections.
// This is a blocking method.
func (s *Server) Run() error {
	// Skip if set to none.
	if s.options.HTTPSocketType != HTTPSocketTypeNone {
		// Set the base glue HTTP handler.
		http.Handle(s.options.HTTPHandleURL, s)

		// Start the http server.
		if s.options.HTTPSocketType == HTTPSocketTypeUnix {
			// Listen on the unix socket.
			l, err := net.Listen("unix", s.options.HTTPListenAddress)
			if err != nil {
				return fmt.Errorf("Listen: %v", err)
			}

			// Start the http server.
			err = http.Serve(l, nil)
			if err != nil {
				return fmt.Errorf("Serve: %v", err)
			}
		} else if s.options.HTTPSocketType == HTTPSocketTypeTCP {
			// Start the http server.
			err := http.ListenAndServe(s.options.HTTPListenAddress, nil)
			if err != nil {
				return fmt.Errorf("ListenAndServe: %v", err)
			}
		} else {
			return fmt.Errorf("invalid socket options type: %v", s.options.HTTPSocketType)
		}
	} else {
		// HINT: This is only a placeholder until the internal glue TCP server is implemented.
		w := make(chan struct{})
		<-w
	}

	return nil
}

// ServeHTTP implements the HTTP Handler interface of the http package.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.bs.ServeHTTP(w, r)
}

//########################//
//### Server - Private ###//
//########################//

func (s *Server) handleOnNewSocketConnection(bs backend.BackendSocket) {
	// Close the socket if incomming connections should be blocked.
	if s.block {
		bs.Close()
		return
	}

	// Create a new socket value.
	// The goroutines are started automatically.
	newSocket(s, bs)
}
<commit_msg>Added GetSocket method to server.<commit_after>/*
 *  Glue - Robust Go and Javascript Socket Library
 *  Copyright (C) 2015  Roland Singer <roland.singer[at]desertbit.com>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

// Package glue - Robust Go and Javascript Socket Library.
// This library is thread-safe.
package glue

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

	"github.com/desertbit/glue/backend"
)

//####################//
//### Public Types ###//
//####################//

// OnNewSocketFunc is an event function.
type OnNewSocketFunc func(s *Socket)

//###################//
//### Server Type ###//
//###################//

// A Server represents a glue server which handles incoming socket connections.
type Server struct {
	bs      *backend.Server
	options *Options

	block       bool
	onNewSocket OnNewSocketFunc

	sockets      map[string]*Socket // A map holding all active current sockets.
	socketsMutex sync.Mutex
}

// NewServer creates a new glue server instance.
// One variadic arguments specifies the server options.
func NewServer(o ...Options) *Server {
	// Get or create the options.
	var options *Options
	if len(o) > 0 {
		options = &o[0]
	} else {
		options = &Options{}
	}

	// Set the default option values for unset values.
	options.SetDefaults()

	// Create a new backend server.
	bs := backend.NewServer(len(options.HTTPHandleURL))

	// Create a new server value.
	s := &Server{
		bs:          bs,
		options:     options,
		onNewSocket: func(*Socket) {}, // Initialize with dummy function to remove nil check.
		sockets:     make(map[string]*Socket),
	}

	// Set the backend server event function.
	bs.OnNewSocketConnection(s.handleOnNewSocketConnection)

	return s
}

// Block new incomming connections.
func (s *Server) Block(b bool) {
	s.block = b
}

// OnNewSocket sets the event function which is
// triggered if a new socket connection was made.
// The event function must not block! As soon as the event function
// returns, the socket is added to the active sockets map.
func (s *Server) OnNewSocket(f OnNewSocketFunc) {
	s.onNewSocket = f
}

// GetSocket obtains a socket by its ID.
// Returns nil if not found.
func (s *Server) GetSocket(id string) *Socket {
	// Lock the mutex.
	s.socketsMutex.Lock()
	defer s.socketsMutex.Unlock()

	// Obtain the socket.
	socket, ok := s.sockets[id]
	if !ok {
		return nil
	}

	return socket
}

// Sockets returns a list of all current connected sockets.
// Hint: Sockets are added to the active sockets list before the OnNewSocket
// event function is called.
// Use the IsInitialized flag to determind if a socket is not ready yet...
func (s *Server) Sockets() []*Socket {
	// Lock the mutex.
	s.socketsMutex.Lock()
	defer s.socketsMutex.Unlock()

	// Create the slice.
	list := make([]*Socket, len(s.sockets))

	// Add all sockets from the map.
	i := 0
	for _, s := range s.sockets {
		list[i] = s
		i++
	}

	return list
}

// Release this package. This will block all new incomming socket connections
// and close all current connected sockets.
func (s *Server) Release() {
	// Block all new incomming socket connections.
	s.Block(true)

	// Wait for a little moment, so new incomming sockets are added
	// to the sockets active list.
	time.Sleep(200 * time.Millisecond)

	// Close all current connected sockets.
	sockets := s.Sockets()
	for _, s := range sockets {
		s.Close()
	}
}

// Run starts the server and listens for incoming socket connections.
// This is a blocking method.
func (s *Server) Run() error {
	// Skip if set to none.
	if s.options.HTTPSocketType != HTTPSocketTypeNone {
		// Set the base glue HTTP handler.
		http.Handle(s.options.HTTPHandleURL, s)

		// Start the http server.
		if s.options.HTTPSocketType == HTTPSocketTypeUnix {
			// Listen on the unix socket.
			l, err := net.Listen("unix", s.options.HTTPListenAddress)
			if err != nil {
				return fmt.Errorf("Listen: %v", err)
			}

			// Start the http server.
			err = http.Serve(l, nil)
			if err != nil {
				return fmt.Errorf("Serve: %v", err)
			}
		} else if s.options.HTTPSocketType == HTTPSocketTypeTCP {
			// Start the http server.
			err := http.ListenAndServe(s.options.HTTPListenAddress, nil)
			if err != nil {
				return fmt.Errorf("ListenAndServe: %v", err)
			}
		} else {
			return fmt.Errorf("invalid socket options type: %v", s.options.HTTPSocketType)
		}
	} else {
		// HINT: This is only a placeholder until the internal glue TCP server is implemented.
		w := make(chan struct{})
		<-w
	}

	return nil
}

// ServeHTTP implements the HTTP Handler interface of the http package.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.bs.ServeHTTP(w, r)
}

//########################//
//### Server - Private ###//
//########################//

func (s *Server) handleOnNewSocketConnection(bs backend.BackendSocket) {
	// Close the socket if incomming connections should be blocked.
	if s.block {
		bs.Close()
		return
	}

	// Create a new socket value.
	// The goroutines are started automatically.
	newSocket(s, bs)
}
<|endoftext|>


<commit_before>package quic

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"os"

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

type PacketHandler interface {
	HandlePacket(addr interface{}, publicHeader *PublicHeader, r *bytes.Reader)
	Run()
}

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

	signer crypto.Signer
	scfg   *handshake.ServerConfig

	sessions map[protocol.ConnectionID]PacketHandler

	streamCallback StreamCallback

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

// NewServer makes a new server
func NewServer(certPath, keyPath string, cb StreamCallback) (*Server, error) {
	path := os.Getenv("GOPATH") + "/src/github.com/lucas-clemente/quic-go/example/"
	signer, err := crypto.NewRSASigner(path+"cert.der", path+"key.der")
	if err != nil {
		return nil, err
	}

	scfg := handshake.NewServerConfig(crypto.NewCurve25519KEX(), signer)

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

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

	s.conn, err = net.ListenUDP("udp", addr)
	if err != nil {
		return err
	}

	for {
		data := make([]byte, protocol.MaxPacketSize)
		n, remoteAddr, err := s.conn.ReadFromUDP(data)
		if err != nil {
			return err
		}
		data = data[:n]
		if err := s.handlePacket(s.conn, remoteAddr, data); err != nil {
			fmt.Printf("error handling packet: %s", err.Error())
		}
	}
}

// Close the server
func (s *Server) Close() error {
	return s.conn.Close()
}

func (s *Server) handlePacket(conn *net.UDPConn, remoteAddr *net.UDPAddr, packet []byte) error {
	r := bytes.NewReader(packet)
	// ToDo: check packet size and send errorcodes.QUIC_PACKET_TOO_LARGE if packet is too large

	publicHeader, err := ParsePublicHeader(r)
	if err != nil {
		// ToDo: send errorcodes.QUIC_INVALID_PACKET_HEADER
		return errors.New("Could not parse public header")
	}
	publicHeader.Raw = packet[:len(packet)-r.Len()]

	// Send Version Negotiation Packet if the client is speaking a different protocol version
	if publicHeader.VersionFlag && !protocol.IsSupportedVersion(publicHeader.VersionNumber) {
		fmt.Println("Sending VersionNegotiationPacket")
		_, err = conn.WriteToUDP(composeVersionNegotiation(publicHeader.ConnectionID), remoteAddr)
		if err != nil {
			return err
		}
		return nil
	}

	session, ok := s.sessions[publicHeader.ConnectionID]
	if !ok {
		fmt.Printf("Serving new connection: %d from %v\n", publicHeader.ConnectionID, remoteAddr)
		session = s.newSession(
			&udpConn{conn: conn, currentAddr: remoteAddr},
			publicHeader.VersionNumber,
			publicHeader.ConnectionID,
			s.scfg,
			s.streamCallback,
		)
		go session.Run()
		s.sessions[publicHeader.ConnectionID] = session
	}
	session.HandlePacket(remoteAddr, publicHeader, r)
	return nil
}

func composeVersionNegotiation(connectionID protocol.ConnectionID) []byte {
	fullReply := &bytes.Buffer{}
	responsePublicHeader := PublicHeader{
		ConnectionID: connectionID,
		PacketNumber: 1,
		VersionFlag:  true,
	}
	err := responsePublicHeader.WritePublicHeader(fullReply)
	if err != nil {
		panic(err) // Should not happen ;)
	}
	fullReply.Write(protocol.SupportedVersionsAsTags)
	return fullReply.Bytes()
}
<commit_msg>use certPath and keyPath variables when setting up the server<commit_after>package quic

import (
	"bytes"
	"errors"
	"fmt"
	"net"

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

type PacketHandler interface {
	HandlePacket(addr interface{}, publicHeader *PublicHeader, r *bytes.Reader)
	Run()
}

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

	signer crypto.Signer
	scfg   *handshake.ServerConfig

	sessions map[protocol.ConnectionID]PacketHandler

	streamCallback StreamCallback

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

// NewServer makes a new server
func NewServer(certPath, keyPath string, cb StreamCallback) (*Server, error) {
	signer, err := crypto.NewRSASigner(certPath, keyPath)
	if err != nil {
		return nil, err
	}

	scfg := handshake.NewServerConfig(crypto.NewCurve25519KEX(), signer)

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

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

	s.conn, err = net.ListenUDP("udp", addr)
	if err != nil {
		return err
	}

	for {
		data := make([]byte, protocol.MaxPacketSize)
		n, remoteAddr, err := s.conn.ReadFromUDP(data)
		if err != nil {
			return err
		}
		data = data[:n]
		if err := s.handlePacket(s.conn, remoteAddr, data); err != nil {
			fmt.Printf("error handling packet: %s", err.Error())
		}
	}
}

// Close the server
func (s *Server) Close() error {
	return s.conn.Close()
}

func (s *Server) handlePacket(conn *net.UDPConn, remoteAddr *net.UDPAddr, packet []byte) error {
	r := bytes.NewReader(packet)
	// ToDo: check packet size and send errorcodes.QUIC_PACKET_TOO_LARGE if packet is too large

	publicHeader, err := ParsePublicHeader(r)
	if err != nil {
		// ToDo: send errorcodes.QUIC_INVALID_PACKET_HEADER
		return errors.New("Could not parse public header")
	}
	publicHeader.Raw = packet[:len(packet)-r.Len()]

	// Send Version Negotiation Packet if the client is speaking a different protocol version
	if publicHeader.VersionFlag && !protocol.IsSupportedVersion(publicHeader.VersionNumber) {
		fmt.Println("Sending VersionNegotiationPacket")
		_, err = conn.WriteToUDP(composeVersionNegotiation(publicHeader.ConnectionID), remoteAddr)
		if err != nil {
			return err
		}
		return nil
	}

	session, ok := s.sessions[publicHeader.ConnectionID]
	if !ok {
		fmt.Printf("Serving new connection: %d from %v\n", publicHeader.ConnectionID, remoteAddr)
		session = s.newSession(
			&udpConn{conn: conn, currentAddr: remoteAddr},
			publicHeader.VersionNumber,
			publicHeader.ConnectionID,
			s.scfg,
			s.streamCallback,
		)
		go session.Run()
		s.sessions[publicHeader.ConnectionID] = session
	}
	session.HandlePacket(remoteAddr, publicHeader, r)
	return nil
}

func composeVersionNegotiation(connectionID protocol.ConnectionID) []byte {
	fullReply := &bytes.Buffer{}
	responsePublicHeader := PublicHeader{
		ConnectionID: connectionID,
		PacketNumber: 1,
		VersionFlag:  true,
	}
	err := responsePublicHeader.WritePublicHeader(fullReply)
	if err != nil {
		panic(err) // Should not happen ;)
	}
	fullReply.Write(protocol.SupportedVersionsAsTags)
	return fullReply.Bytes()
}
<|endoftext|>


<commit_before>package tao

import (
  "log"
  "net"
  "os"
  "time"
  "sync"
  "crypto/tls"
  "crypto/rand"
)

func init() {
  netIdentifier = NewAtomicInt64(0)
  tlsWrapper = func(conn net.Conn) net.Conn {
    return conn
  }
}

var (
  netIdentifier *AtomicInt64
  tlsWrapper func(net.Conn) net.Conn
)

type Server interface{
  IsRunning() bool
  GetAllConnections() *ConcurrentMap
  GetTimingWheel() *TimingWheel
  GetWorkerPool() *WorkerPool
  GetServerAddress() string
  Start()
  Close()

  SetOnScheduleCallback(time.Duration, func(time.Time, interface{}))
  GetOnScheduleCallback() (time.Duration, onScheduleFunc)
  SetOnConnectCallback(func(Connection)bool)
  GetOnConnectCallback() onConnectFunc
  SetOnMessageCallback(func(Message, Connection))
  GetOnMessageCallback() onMessageFunc
  SetOnCloseCallback(func(Connection))
  GetOnCloseCallback() onCloseFunc
  SetOnErrorCallback(func())
  GetOnErrorCallback() onErrorFunc
}

type TCPServer struct{
  isRunning *AtomicBoolean
  connections *ConcurrentMap
  timingWheel *TimingWheel
  workerPool *WorkerPool
  finish *sync.WaitGroup
  address string
  closeServChan chan struct{}

  onConnect onConnectFunc
  onMessage onMessageFunc
  onClose onCloseFunc
  onError onErrorFunc

  duration time.Duration
  onSchedule onScheduleFunc
}

func NewTCPServer(addr string) Server {
  server := &TCPServer {
    isRunning: NewAtomicBoolean(true),
    connections: NewConcurrentMap(),
    timingWheel: NewTimingWheel(),
    workerPool: NewWorkerPool(WORKERS),
    finish: &sync.WaitGroup{},
    address: addr,
  }
  server.finish.Add(1)
  go server.timeOutLoop()
  return server
}

func (server *TCPServer)IsRunning() bool {
  return server.isRunning.Get()
}

func (server *TCPServer)GetAllConnections() *ConcurrentMap {
  return server.connections
}

func (server *TCPServer)GetTimingWheel() *TimingWheel {
  return server.timingWheel
}

func (server *TCPServer)GetWorkerPool() *WorkerPool {
  return server.workerPool
}

func (server *TCPServer)GetServerAddress() string {
  return server.address
}

func (server *TCPServer) Start() {
  listener, err := net.Listen("tcp", server.address)
  if err != nil {
    log.Fatalln(err)
  }
  defer listener.Close()

  for server.IsRunning() {
    conn, err := listener.Accept()
    if err != nil {
      log.Fatalln(err)
    }

    conn = tlsWrapper(conn)  // wrap as a tls connection if configured

    /* Create a TCP connection upon accepting a new client, assign an net id
    to it, then manage it in connections map, and start it */
    netid := netIdentifier.GetAndIncrement()
    tcpConn := NewServerConnection(netid, server, conn)
    tcpConn.SetName(tcpConn.GetRemoteAddress().String())
    if server.connections.Size() < MAX_CONNECTIONS {
      duration, onSchedule := server.GetOnScheduleCallback()
      if onSchedule != nil {
        tcpConn.RunEvery(duration, onSchedule)
      }
      server.connections.Put(netid, tcpConn)

      // put tcpConn.Start() run in another WaitGroup-synchronized go-routine
      server.finish.Add(1)
      go func() {
        tcpConn.Start()
      }()

      log.Printf("Accepting client %s, net id %d, now %d\n", tcpConn.GetName(), netid, server.connections.Size())
      for v := range server.connections.IterValues() {
        log.Printf("Client %s\n", v.(Connection).GetName())
      }
    } else {
      log.Printf("WARN, MAX CONNS %d, refuse\n", MAX_CONNECTIONS)
      tcpConn.Close()
    }
  }
}

// wait until all connections closed
func (server *TCPServer) Close() {
  if server.isRunning.CompareAndSet(true, false) {
    for v := range server.GetAllConnections().IterValues() {
      c := v.(Connection)
      c.Close()
    }
    close(server.closeServChan)
    server.finish.Wait()
    os.Exit(0)
  }
}

func (server *TCPServer)GetTimeOutChannel() chan *OnTimeOut {
  return server.timingWheel.TimeOutChan
}

func (server *TCPServer)SetOnScheduleCallback(duration time.Duration, callback func(time.Time, interface{})) {
  server.duration = duration
  server.onSchedule = onScheduleFunc(callback)
}

func (server *TCPServer)GetOnScheduleCallback() (time.Duration, onScheduleFunc) {
  return server.duration, server.onSchedule
}

func (server *TCPServer)SetOnConnectCallback(callback func(Connection)bool) {
  server.onConnect = onConnectFunc(callback)
}

func (server *TCPServer)GetOnConnectCallback() onConnectFunc {
  return server.onConnect
}

func (server *TCPServer)SetOnMessageCallback(callback func(Message, Connection)) {
  server.onMessage = onMessageFunc(callback)
}

func (server *TCPServer)GetOnMessageCallback() onMessageFunc {
  return server.onMessage
}

func (server *TCPServer)SetOnCloseCallback(callback func(Connection)) {
  server.onClose = onCloseFunc(callback)
}

func (server *TCPServer)GetOnCloseCallback() onCloseFunc {
  return server.onClose
}

func (server *TCPServer)SetOnErrorCallback(callback func()) {
  server.onError = onErrorFunc(callback)
}

func (server *TCPServer)GetOnErrorCallback() onErrorFunc {
  return server.onError
}

type TLSTCPServer struct{
  certFile string
  keyFile string
  *TCPServer
}

func NewTLSTCPServer(addr, cert, key string) Server {
  server := &TLSTCPServer {
    certFile: cert,
    keyFile: key,
    TCPServer: NewTCPServer(addr).(*TCPServer),
  }

  config, err := loadTLSConfig(server.certFile, server.keyFile)
  if err != nil {
    log.Fatalln(err)
  }

  setTLSWrapper(func(conn net.Conn) net.Conn{
    return tls.Server(conn, &config)
  })

  return server
}

func (server *TLSTCPServer)IsRunning() bool {
  return server.TCPServer.IsRunning()
}

func (server *TLSTCPServer)GetAllConnections() *ConcurrentMap {
  return server.TCPServer.GetAllConnections()
}

func (server *TLSTCPServer)GetTimingWheel() *TimingWheel {
  return server.TCPServer.GetTimingWheel()
}

func (server *TLSTCPServer)GetWorkerPool() *WorkerPool {
  return server.TCPServer.GetWorkerPool()
}

func (server *TLSTCPServer)GetServerAddress() string {
  return server.TCPServer.GetServerAddress()
}

func (server *TLSTCPServer)Start() {
  server.TCPServer.Start()
}

func (server *TLSTCPServer)Close() {
  server.TCPServer.Close()
}

func (server *TLSTCPServer)SetOnScheduleCallback(duration time.Duration, callback func(time.Time, interface{})) {
  server.TCPServer.SetOnScheduleCallback(duration, callback)
}

func (server *TLSTCPServer)GetOnScheduleCallback() (time.Duration, onScheduleFunc) {
  return server.TCPServer.GetOnScheduleCallback()
}

func (server *TLSTCPServer)SetOnConnectCallback(callback func(Connection)bool) {
  server.TCPServer.SetOnConnectCallback(callback)
}

func (server *TLSTCPServer)GetOnConnectCallback() onConnectFunc {
  return server.TCPServer.GetOnConnectCallback()
}

func (server *TLSTCPServer)SetOnMessageCallback(callback func(Message, Connection)) {
  server.TCPServer.SetOnMessageCallback(callback)
}

func (server *TLSTCPServer)GetOnMessageCallback() onMessageFunc {
  return server.TCPServer.GetOnMessageCallback()
}

func (server *TLSTCPServer)SetOnCloseCallback(callback func(Connection)) {
  server.TCPServer.SetOnCloseCallback(callback)
}

func (server *TLSTCPServer)GetOnCloseCallback() onCloseFunc {
  return server.TCPServer.GetOnCloseCallback()
}

func (server *TLSTCPServer)SetOnErrorCallback(callback func()) {
  server.TCPServer.SetOnErrorCallback(callback)
}

func (server *TLSTCPServer)GetOnErrorCallback() onErrorFunc {
  return server.TCPServer.GetOnErrorCallback()
}

/* Retrieve the extra data(i.e. net id), and then redispatch
timeout callbacks to corresponding client connection, this
prevents one client from running callbacks of other clients */

func (server *TCPServer) timeOutLoop() {
  defer server.finish.Done()

  for {
    select {
    case <-server.closeServChan:
      return

    case timeout := <-server.GetTimingWheel().TimeOutChan:
      netid := timeout.ExtraData.(int64)
      if conn, ok := server.connections.Get(netid); ok {
        tcpConn := conn.(Connection)
        if !tcpConn.IsClosed() {
          tcpConn.GetTimeOutChannel()<- timeout
        }
      } else {
        log.Printf("Invalid client %d\n", netid)
      }
    }
  }
}

func loadTLSConfig(certFile, keyFile string) (tls.Config, error) {
  var config tls.Config
  cert, err := tls.LoadX509KeyPair(certFile, keyFile)
  return config, err
  config = tls.Config{Certificates: []tls.Certificate{cert}}
  now := time.Now()
  config.Time = func() time.Time { return now }
  config.Rand = rand.Reader
  return config, nil
}

func setTLSWrapper(wrapper func(conn net.Conn) net.Conn) {
  tlsWrapper = wrapper
}
<commit_msg>activate time out loop in TCPServer::Start()<commit_after>  package tao

import (
  "log"
  "net"
  "os"
  "time"
  "sync"
  "crypto/tls"
  "crypto/rand"
)

func init() {
  netIdentifier = NewAtomicInt64(0)
  tlsWrapper = func(conn net.Conn) net.Conn {
    return conn
  }
}

var (
  netIdentifier *AtomicInt64
  tlsWrapper func(net.Conn) net.Conn
)

type Server interface{
  IsRunning() bool
  GetAllConnections() *ConcurrentMap
  GetTimingWheel() *TimingWheel
  GetWorkerPool() *WorkerPool
  GetServerAddress() string
  Start()
  Close()

  SetOnScheduleCallback(time.Duration, func(time.Time, interface{}))
  GetOnScheduleCallback() (time.Duration, onScheduleFunc)
  SetOnConnectCallback(func(Connection)bool)
  GetOnConnectCallback() onConnectFunc
  SetOnMessageCallback(func(Message, Connection))
  GetOnMessageCallback() onMessageFunc
  SetOnCloseCallback(func(Connection))
  GetOnCloseCallback() onCloseFunc
  SetOnErrorCallback(func())
  GetOnErrorCallback() onErrorFunc
}

type TCPServer struct{
  isRunning *AtomicBoolean
  connections *ConcurrentMap
  timingWheel *TimingWheel
  workerPool *WorkerPool
  finish *sync.WaitGroup
  address string
  closeServChan chan struct{}

  onConnect onConnectFunc
  onMessage onMessageFunc
  onClose onCloseFunc
  onError onErrorFunc

  duration time.Duration
  onSchedule onScheduleFunc
}

func NewTCPServer(addr string) Server {
  return &TCPServer {
    isRunning: NewAtomicBoolean(true),
    connections: NewConcurrentMap(),
    timingWheel: NewTimingWheel(),
    workerPool: NewWorkerPool(WORKERS),
    finish: &sync.WaitGroup{},
    address: addr,
  }
}

func (server *TCPServer)IsRunning() bool {
  return server.isRunning.Get()
}

func (server *TCPServer)GetAllConnections() *ConcurrentMap {
  return server.connections
}

func (server *TCPServer)GetTimingWheel() *TimingWheel {
  return server.timingWheel
}

func (server *TCPServer)GetWorkerPool() *WorkerPool {
  return server.workerPool
}

func (server *TCPServer)GetServerAddress() string {
  return server.address
}

func (server *TCPServer) Start() {
  server.finish.Add(1)
  go server.timeOutLoop()

  listener, err := net.Listen("tcp", server.address)
  if err != nil {
    log.Fatalln(err)
  }
  defer listener.Close()

  for server.IsRunning() {
    conn, err := listener.Accept()
    if err != nil {
      log.Println("Accept error - ", err)
      continue
    }

    conn = tlsWrapper(conn)  // wrap as a tls connection if configured

    /* Create a TCP connection upon accepting a new client, assign an net id
    to it, then manage it in connections map, and start it */
    netid := netIdentifier.GetAndIncrement()
    tcpConn := NewServerConnection(netid, server, conn)
    tcpConn.SetName(tcpConn.GetRemoteAddress().String())
    if server.connections.Size() < MAX_CONNECTIONS {
      duration, onSchedule := server.GetOnScheduleCallback()
      if onSchedule != nil {
        tcpConn.RunEvery(duration, onSchedule)
      }
      server.connections.Put(netid, tcpConn)

      // put tcpConn.Start() run in another WaitGroup-synchronized go-routine
      server.finish.Add(1)
      go func() {
        tcpConn.Start()
      }()

      log.Printf("Accepting client %s, net id %d, now %d\n", tcpConn.GetName(), netid, server.connections.Size())
      for v := range server.connections.IterValues() {
        log.Printf("Client %s\n", v.(Connection).GetName())
      }
    } else {
      log.Printf("WARN, MAX CONNS %d, refuse\n", MAX_CONNECTIONS)
      tcpConn.Close()
    }
  }
}

// wait until all connections closed
func (server *TCPServer) Close() {
  if server.isRunning.CompareAndSet(true, false) {
    for v := range server.GetAllConnections().IterValues() {
      c := v.(Connection)
      c.Close()
    }
    close(server.closeServChan)
    server.finish.Wait()
    os.Exit(0)
  }
}

func (server *TCPServer)GetTimeOutChannel() chan *OnTimeOut {
  return server.timingWheel.TimeOutChan
}

func (server *TCPServer)SetOnScheduleCallback(duration time.Duration, callback func(time.Time, interface{})) {
  server.duration = duration
  server.onSchedule = onScheduleFunc(callback)
}

func (server *TCPServer)GetOnScheduleCallback() (time.Duration, onScheduleFunc) {
  return server.duration, server.onSchedule
}

func (server *TCPServer)SetOnConnectCallback(callback func(Connection)bool) {
  server.onConnect = onConnectFunc(callback)
}

func (server *TCPServer)GetOnConnectCallback() onConnectFunc {
  return server.onConnect
}

func (server *TCPServer)SetOnMessageCallback(callback func(Message, Connection)) {
  server.onMessage = onMessageFunc(callback)
}

func (server *TCPServer)GetOnMessageCallback() onMessageFunc {
  return server.onMessage
}

func (server *TCPServer)SetOnCloseCallback(callback func(Connection)) {
  server.onClose = onCloseFunc(callback)
}

func (server *TCPServer)GetOnCloseCallback() onCloseFunc {
  return server.onClose
}

func (server *TCPServer)SetOnErrorCallback(callback func()) {
  server.onError = onErrorFunc(callback)
}

func (server *TCPServer)GetOnErrorCallback() onErrorFunc {
  return server.onError
}

type TLSTCPServer struct{
  certFile string
  keyFile string
  *TCPServer
}

func NewTLSTCPServer(addr, cert, key string) Server {
  server := &TLSTCPServer {
    certFile: cert,
    keyFile: key,
    TCPServer: NewTCPServer(addr).(*TCPServer),
  }

  config, err := loadTLSConfig(server.certFile, server.keyFile)
  if err != nil {
    log.Fatalln(err)
  }

  setTLSWrapper(func(conn net.Conn) net.Conn{
    return tls.Server(conn, &config)
  })

  return server
}

func (server *TLSTCPServer)IsRunning() bool {
  return server.TCPServer.IsRunning()
}

func (server *TLSTCPServer)GetAllConnections() *ConcurrentMap {
  return server.TCPServer.GetAllConnections()
}

func (server *TLSTCPServer)GetTimingWheel() *TimingWheel {
  return server.TCPServer.GetTimingWheel()
}

func (server *TLSTCPServer)GetWorkerPool() *WorkerPool {
  return server.TCPServer.GetWorkerPool()
}

func (server *TLSTCPServer)GetServerAddress() string {
  return server.TCPServer.GetServerAddress()
}

func (server *TLSTCPServer)Start() {
  server.TCPServer.Start()
}

func (server *TLSTCPServer)Close() {
  server.TCPServer.Close()
}

func (server *TLSTCPServer)SetOnScheduleCallback(duration time.Duration, callback func(time.Time, interface{})) {
  server.TCPServer.SetOnScheduleCallback(duration, callback)
}

func (server *TLSTCPServer)GetOnScheduleCallback() (time.Duration, onScheduleFunc) {
  return server.TCPServer.GetOnScheduleCallback()
}

func (server *TLSTCPServer)SetOnConnectCallback(callback func(Connection)bool) {
  server.TCPServer.SetOnConnectCallback(callback)
}

func (server *TLSTCPServer)GetOnConnectCallback() onConnectFunc {
  return server.TCPServer.GetOnConnectCallback()
}

func (server *TLSTCPServer)SetOnMessageCallback(callback func(Message, Connection)) {
  server.TCPServer.SetOnMessageCallback(callback)
}

func (server *TLSTCPServer)GetOnMessageCallback() onMessageFunc {
  return server.TCPServer.GetOnMessageCallback()
}

func (server *TLSTCPServer)SetOnCloseCallback(callback func(Connection)) {
  server.TCPServer.SetOnCloseCallback(callback)
}

func (server *TLSTCPServer)GetOnCloseCallback() onCloseFunc {
  return server.TCPServer.GetOnCloseCallback()
}

func (server *TLSTCPServer)SetOnErrorCallback(callback func()) {
  server.TCPServer.SetOnErrorCallback(callback)
}

func (server *TLSTCPServer)GetOnErrorCallback() onErrorFunc {
  return server.TCPServer.GetOnErrorCallback()
}

/* Retrieve the extra data(i.e. net id), and then redispatch
timeout callbacks to corresponding client connection, this
prevents one client from running callbacks of other clients */

func (server *TCPServer) timeOutLoop() {
  defer server.finish.Done()

  for {
    select {
    case <-server.closeServChan:
      return

    case timeout := <-server.GetTimingWheel().TimeOutChan:
      netid := timeout.ExtraData.(int64)
      if conn, ok := server.connections.Get(netid); ok {
        tcpConn := conn.(Connection)
        if !tcpConn.IsClosed() {
          tcpConn.GetTimeOutChannel()<- timeout
        }
      } else {
        log.Printf("Invalid client %d\n", netid)
      }
    }
  }
}

func loadTLSConfig(certFile, keyFile string) (tls.Config, error) {
  var config tls.Config
  cert, err := tls.LoadX509KeyPair(certFile, keyFile)
  return config, err
  config = tls.Config{Certificates: []tls.Certificate{cert}}
  now := time.Now()
  config.Time = func() time.Time { return now }
  config.Rand = rand.Reader
  return config, nil
}

func setTLSWrapper(wrapper func(conn net.Conn) net.Conn) {
  tlsWrapper = wrapper
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base32"
	"encoding/hex"
	"fmt"
	"io"

	"github.com/dgryski/dgoogauth"
	"github.com/gin-gonic/contrib/sessions"
	"github.com/gin-gonic/gin"
	qrcode "github.com/skip2/go-qrcode"
)

const (
	OTPConfigWindowSize = 10
	OTPConfigUTC        = true
)

// JSONResponse ajax response
type JSONResponse struct {
	Ok  bool   `json:"ok"`
	Txt string `json:"txt"`
}

// generate a new key
func generateSecret() (secret string, err error) {
	var key []byte
	key = make([]byte, sha1.Size)
	_, err = io.ReadFull(rand.Reader, key)
	if err != nil {
		return
	}
	secret = base32.StdEncoding.EncodeToString(key)
	return
}

// Step 1: init OTPconfig en generate QRRCode
func handlergenQrCode(c *gin.Context) {
	var err error
	data := struct {
		user string `binding:"required"`
	}{}
	if err = c.Bind(&data); err != nil {
		return
	}

	// OTPConfig
	secret, err := generateSecret()
	if err != nil {
		return
	}

	config := dgoogauth.OTPConfig{
		Secret:     secret,
		WindowSize: OTPConfigWindowSize,
		UTC:        OTPConfigUTC,
	}

	// gen qrcode
	hasher := md5.New()
	hasher.Write([]byte(data.user))
	filename := hex.EncodeToString(hasher.Sum(nil)) + ".png"
	// gen url
	url := config.ProvisionURIWithIssuer(data.user, "otpDemo")
	// gen & save qrcode
	err = qrcode.WriteFile(url, qrcode.Medium, 256, "./public/qrcodes/"+filename)
	// save OTPConfig in session (don't do this at /home !)
	session := sessions.Default(c)
	session.Set("optconfig", config)
	session.Save()
	c.JSON(200, JSONResponse{true, filename})
}

// step 3 check code
func handlerCheckCode(c *gin.Context) {
	var err error
	data := struct {
		code string `binding:"required"`
	}{}
	if err = c.Bind(&data); err != nil {
		return
	}
	session := sessions.Default(c)
	configInt := session.Get("optconfig")
	if configInt == nil {
		c.JSON(400, JSONResponse{false, "no OTPConfing found on session"})
		return
	}
	config := configInt.(dgoogauth.OTPConfig)
	valid, err := config.Authenticate(data.code)
	if err != nil {
		c.JSON(400, JSONResponse{false, "no OTPConfing found on session"})
		return
	}
	c.JSON(200, JSONResponse{true, fmt.Sprintf("%t", valid)})
}

func main() {
	r := gin.Default()
	store := sessions.NewCookieStore([]byte("bigsecret"))
	r.Use(sessions.Sessions("otp", store))
	r.Static("/public", "./public")
	// index
	r.StaticFile("/", "./public/index.html")

	// gen qrcode
	r.POST("/aj/genQRCode", handlergenQrCode)

	// check code
	r.POST("/aj/checkCode", handlerCheckCode)

	//r.GET("/", handlerIndex)
	r.Run() // listen and server on 0.0.0.0:8080
}
<commit_msg> test on secret size<commit_after>package main

import (
	"crypto/md5"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base32"
	"encoding/gob"
	"encoding/hex"
	"fmt"
	"log"

	"github.com/dgryski/dgoogauth"
	"github.com/gin-gonic/contrib/sessions"
	"github.com/gin-gonic/gin"
	qrcode "github.com/skip2/go-qrcode"
)

const (
	OTPConfigWindowSize = 10
	OTPConfigUTC        = true
)

// JSONResponse ajax response
type JSONResponse struct {
	Ok  bool   `json:"ok"`
	Txt string `json:"txt"`
}

// generate a new key
func generateSecret() (secret string, err error) {
	var key []byte
	key = make([]byte, sha1.Size)
	_, err = rand.Read(key)
	if err != nil {
		return
	}
	secret = base32.StdEncoding.EncodeToString(key)
	return
}

// Step 1: init OTPconfig en generate QRRCode
func handlergenQrCode(c *gin.Context) {
	var err error
	data := struct {
		user string `binding:"required"`
	}{}
	if err = c.Bind(&data); err != nil {
		return
	}

	// OTPConfig
	secret, err := generateSecret()
	if err != nil {
		return
	}

	fmt.Println(secret)

	config := dgoogauth.OTPConfig{
		Secret:     secret,
		WindowSize: OTPConfigWindowSize,
		UTC:        OTPConfigUTC,
	}

	// gen qrcode
	hasher := md5.New()
	hasher.Write([]byte(data.user))
	filename := hex.EncodeToString(hasher.Sum(nil)) + ".png"
	// gen url
	url := config.ProvisionURIWithIssuer(data.user, "otpDemo")
	log.Println(url)
	// gen & save qrcode
	err = qrcode.WriteFile(url, qrcode.Medium, 256, "./public/qrcodes/"+filename)
	// save OTPConfig in session (don't do this at /home !)
	session := sessions.Default(c)
	session.Set("optconfig", config)
	err = session.Save()
	fmt.Println(err)

	c.JSON(200, JSONResponse{true, filename})
}

// step 3 check code
func handlerCheckCode(c *gin.Context) {
	var err error
	data := struct {
		code string `binding:"required"`
	}{}
	if err = c.Bind(&data); err != nil {
		return
	}
	session := sessions.Default(c)
	configInt := session.Get("optconfig")
	if configInt == nil {
		c.JSON(400, JSONResponse{false, "no OTPConfing found on session"})
		return
	}
	config := configInt.(dgoogauth.OTPConfig)
	log.Println(configInt)
	valid, err := config.Authenticate(data.code)
	fmt.Println(err)
	if err != nil && err != dgoogauth.ErrInvalidCode {
		c.JSON(400, JSONResponse{false, err.Error()})
		return
	}
	c.JSON(200, JSONResponse{true, fmt.Sprintf("%t", valid)})
}

func main() {
	gob.Register(dgoogauth.OTPConfig{})
	r := gin.Default()
	store := sessions.NewCookieStore([]byte("bigsecret"))
	r.Use(sessions.Sessions("otp", store))
	r.Static("/public", "./public")
	// index
	r.StaticFile("/", "./public/index.html")

	// gen qrcode
	r.POST("/aj/genQRCode", handlergenQrCode)

	// check code
	r.POST("/aj/checkCode", handlerCheckCode)

	//r.GET("/", handlerIndex)
	r.Run() // listen and server on 0.0.0.0:8080
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"flag"
	"github.com/nfnt/resize"
	"image"
	"image/color"
	_ "image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"io/ioutil"
	"launchpad.net/goyaml"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

type proxyHandler struct {
	Host      string // host of proxy
	Scheme    string // scheme of the proxy (typically http)
	Prefix    string // path prefix for the proxy
	Format    string // image format on the proxy
	Watermark string // filename of the watermark

	watermarkImage  image.Image
	backgroundColor color.Color
}

func (h proxyHandler) urlFor(u *url.URL) *url.URL {
	n := new(url.URL)
	n.Scheme = h.Scheme
	n.Host = h.Host

  if (len(strings.Split(u.Path, "/")) == 3) {
    n.Path = h.Prefix + "/" + strings.Join(strings.Split(u.Path, "/")[2:], "/")
  } else {
    n.Path = h.Prefix + u.Path
  }

  log.Println(n.Path)

	n.Path = n.Path[0:strings.LastIndex(n.Path, ".")]
	n.Path = n.Path + "." + h.Format
	// n.Path = "/photos/jpg/022982fd1b29cadfc35588422e15d380f13dc8fe.jpg"
	n.RawQuery = ""
	return n
}

func (h *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// TODO: Recover from panic and throw a 500
	var (
		status           int
		duration         time.Duration
		err1, err2, err3 error
		resp             *http.Response
		m                image.Image
	)
	startTime := time.Now()

	resp, err1 = http.Get(h.urlFor(r.URL).String())
	defer resp.Body.Close()

	if err1 != nil || resp.StatusCode != http.StatusOK {
		status = resp.StatusCode
		w.WriteHeader(status)
	}

	if err1 == nil && resp.StatusCode == http.StatusOK {
		m, _, err2 = image.Decode(resp.Body)
		if err2 != nil {
			status = http.StatusInternalServerError
			w.WriteHeader(status)
		} else {
			status = http.StatusOK
			err3 = h.serveImage(w, r, m)
			if err3 != nil {
				status = http.StatusInternalServerError
				w.WriteHeader(status)
			}
		}
	}

	duration = time.Since(startTime)

	logStatement := "%v\t%v\n"
	logStatement += "                    Completed %d %v in %dms\n"
	if err1 != nil {
		logStatement += "                    error > " + err1.Error()
	}
	if err2 != nil {
		logStatement += "                    error > " + err2.Error()
	}
	if err3 != nil {
		logStatement += "                    error > " + err3.Error()
	}
	log.Printf(logStatement, r.Method, r.URL.String(), status, http.StatusText(status), int(duration.Nanoseconds()/1000000))
}

func (h *proxyHandler) resizeImage(m image.Image, options map[string]interface{}) image.Image {
	dstRect := image.Rect(0, 0, options["width"].(int), options["height"].(int))

	if options["crop"].(bool) {
		m = resizeAndCrop(m, dstRect)
	} else if options["enforce"].(bool) {
		m = resize.Resize(uint(dstRect.Dx()), uint(dstRect.Dy()), m, resize.MitchellNetravali)
	} else if options["minimum"].(bool) {
		srcRect := m.Bounds()
		srcRatio := float32(srcRect.Dx()) / float32(srcRect.Dy())
		dstRatio := float32(dstRect.Dx()) / float32(dstRect.Dy())
		if srcRatio > dstRatio {
			dstRect.Max.X = int(float32(dstRect.Max.Y) * srcRatio)
		} else {
			dstRect.Max.Y = int(float32(dstRect.Max.X) / srcRatio)
		}

		m = resize.Resize(uint(dstRect.Dx()), uint(dstRect.Dy()), m, resize.MitchellNetravali)
	} else {
		m = resizeAndPad(m, dstRect, h.backgroundColor)
	}

	return m
}

func (h *proxyHandler) serveImage(w http.ResponseWriter, r *http.Request, m image.Image) error {
	options := extractOptions(r, m)
	m = h.resizeImage(m, options)

	var err error
	m, err = h.watermark(m)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return err
	}

	const (
		year       = time.Hour * 24 * 7 * 52
		timeFormat = "Mon, 02 Jan 2006 15:04:05 GMT"
	)
	w.Header().Set("Content-type", options["mimeType"].(string))
	w.Header().Set("Expires", time.Now().Add(year).UTC().Format(timeFormat))
	w.Header().Set("Cache-Control", "public, max-age=15724800")

	etag := md5.New()
	body := new(bytes.Buffer)
	bodyWriter := io.MultiWriter(body, etag)

	switch options["mimeType"].(string) {
	case "image/jpeg":
		err = jpeg.Encode(bodyWriter, m, &jpeg.Options{Quality: 85})
	case "image/png":
		err = png.Encode(bodyWriter, m)
	}

	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return err
	} else {
		w.Header().Set("ETag", hex.EncodeToString(etag.Sum(nil)))
		body.WriteTo(w)
	}

	return nil
}

func parseConfigFile(n string) (*proxyHandler, error) {
	file, err := os.Open(n)

	if err != nil {
		return nil, err
	}

	var h *proxyHandler
	data, err := ioutil.ReadAll(file)

	if err != nil {
		return nil, err
	}

	err = goyaml.Unmarshal(data, &h)
	if err != nil {
		return nil, err
	}

  log.Println(h.Watermark)
  if (h.Watermark != "") {
    file, err = os.Open(h.Watermark)
    h.watermarkImage, _, err = image.Decode(file)
    if err != nil {
      return nil, err
    }
  }

	h.backgroundColor = color.RGBA{80, 80, 80, 1}
	return h, nil
}

var configFile string
var serverPort string
var serverBinding string

func init() {
	flag.StringVar(&configFile, "config", "", "conf file (see config.yml.sample)")
	flag.StringVar(&serverPort, "port", "3000", "run the server on the specified port")
	flag.StringVar(&serverBinding, "binding", "0.0.0.0", "bind the server to the specified ip")
}

func main() {
	flag.Parse()
	h, err := parseConfigFile(configFile)
	if err != nil {
		log.Fatal("Unable to find / parse config file.")
	}
	http.Handle("/", h)
	http.ListenAndServe(serverBinding+":"+serverPort, nil)
}
<commit_msg>New config options<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"flag"
	"github.com/nfnt/resize"
	"image"
	"image/color"
	_ "image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"io/ioutil"
	"launchpad.net/goyaml"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

type proxyHandler struct {
	Host      string // host of proxy
	Scheme    string // scheme of the proxy (typically http)
	Port      string // The port the server runs on
	Prefix    string // path prefix for the proxy
	Format    string // image format on the proxy
	Watermark string // filename of the watermark

	watermarkImage  image.Image
	backgroundColor color.Color
}

func (h proxyHandler) urlFor(u *url.URL) *url.URL {
	n := new(url.URL)
	n.Scheme = h.Scheme
	n.Host = h.Host

  if (len(strings.Split(u.Path, "/")) == 3) {
    n.Path = h.Prefix + "/" + strings.Join(strings.Split(u.Path, "/")[2:], "/")
  } else {
    n.Path = h.Prefix + u.Path
  }

  log.Println(n.Path)

	n.Path = n.Path[0:strings.LastIndex(n.Path, ".")]
	n.Path = n.Path + "." + h.Format
	// n.Path = "/photos/jpg/022982fd1b29cadfc35588422e15d380f13dc8fe.jpg"
	n.RawQuery = ""
	return n
}

func (h *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// TODO: Recover from panic and throw a 500
	var (
		status           int
		duration         time.Duration
		err1, err2, err3 error
		resp             *http.Response
		m                image.Image
	)
	startTime := time.Now()

	resp, err1 = http.Get(h.urlFor(r.URL).String())
	defer resp.Body.Close()

	if err1 != nil || resp.StatusCode != http.StatusOK {
		status = resp.StatusCode
		w.WriteHeader(status)
	}

	if err1 == nil && resp.StatusCode == http.StatusOK {
		m, _, err2 = image.Decode(resp.Body)
		if err2 != nil {
			status = http.StatusInternalServerError
			w.WriteHeader(status)
		} else {
			status = http.StatusOK
			err3 = h.serveImage(w, r, m)
			if err3 != nil {
				status = http.StatusInternalServerError
				w.WriteHeader(status)
			}
		}
	}

	duration = time.Since(startTime)

	logStatement := "%v\t%v\n"
	logStatement += "                    Completed %d %v in %dms\n"
	if err1 != nil {
		logStatement += "                    error > " + err1.Error()
	}
	if err2 != nil {
		logStatement += "                    error > " + err2.Error()
	}
	if err3 != nil {
		logStatement += "                    error > " + err3.Error()
	}
	log.Printf(logStatement, r.Method, r.URL.String(), status, http.StatusText(status), int(duration.Nanoseconds()/1000000))
}

func (h *proxyHandler) resizeImage(m image.Image, options map[string]interface{}) image.Image {
	dstRect := image.Rect(0, 0, options["width"].(int), options["height"].(int))

	if options["crop"].(bool) {
		m = resizeAndCrop(m, dstRect)
	} else if options["enforce"].(bool) {
		m = resize.Resize(uint(dstRect.Dx()), uint(dstRect.Dy()), m, resize.MitchellNetravali)
	} else if options["minimum"].(bool) {
		srcRect := m.Bounds()
		srcRatio := float32(srcRect.Dx()) / float32(srcRect.Dy())
		dstRatio := float32(dstRect.Dx()) / float32(dstRect.Dy())
		if srcRatio > dstRatio {
			dstRect.Max.X = int(float32(dstRect.Max.Y) * srcRatio)
		} else {
			dstRect.Max.Y = int(float32(dstRect.Max.X) / srcRatio)
		}

		m = resize.Resize(uint(dstRect.Dx()), uint(dstRect.Dy()), m, resize.MitchellNetravali)
	} else {
		m = resizeAndPad(m, dstRect, h.backgroundColor)
	}

	return m
}

func (h *proxyHandler) serveImage(w http.ResponseWriter, r *http.Request, m image.Image) error {
	options := extractOptions(r, m)
	m = h.resizeImage(m, options)

	var err error
	m, err = h.watermark(m)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return err
	}

	const (
		year       = time.Hour * 24 * 7 * 52
		timeFormat = "Mon, 02 Jan 2006 15:04:05 GMT"
	)
	w.Header().Set("Content-type", options["mimeType"].(string))
	w.Header().Set("Expires", time.Now().Add(year).UTC().Format(timeFormat))
	w.Header().Set("Cache-Control", "public, max-age=15724800")

	etag := md5.New()
	body := new(bytes.Buffer)
	bodyWriter := io.MultiWriter(body, etag)

	switch options["mimeType"].(string) {
	case "image/jpeg":
		err = jpeg.Encode(bodyWriter, m, &jpeg.Options{Quality: 85})
	case "image/png":
		err = png.Encode(bodyWriter, m)
	}

	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return err
	} else {
		w.Header().Set("ETag", hex.EncodeToString(etag.Sum(nil)))
		body.WriteTo(w)
	}

	return nil
}

func parseConfigFile(n string) (*proxyHandler, error) {
	file, err := os.Open(n)

	if err != nil {
		return nil, err
	}

	var h *proxyHandler
	data, err := ioutil.ReadAll(file)

	if err != nil {
		return nil, err
	}

	err = goyaml.Unmarshal(data, &h)
	if err != nil {
		return nil, err
	}

	if (h.Watermark != "") {
		file, err = os.Open(h.Watermark)
		h.watermarkImage, _, err = image.Decode(file)
		if err != nil {
			return nil, err
		}
	}

	h.backgroundColor = color.RGBA{80, 80, 80, 1}
	return h, nil
}

var configFile string
var serverPort string
var serverBinding string

func init() {
	flag.StringVar(&configFile, "config", "", "conf file (see config.yml.sample)")
	flag.StringVar(&serverPort, "port", "3000", "run the server on the specified port")
	flag.StringVar(&serverBinding, "binding", "0.0.0.0", "bind the server to the specified ip")
}

func main() {
	flag.Parse()
	h, err := parseConfigFile(configFile)

	if (h.Port != "") {
		serverPort = h.Port;
	}

	if err != nil {
		log.Fatal("Unable to find / parse config file.")
	}

	http.Handle("/", h)
	http.ListenAndServe(serverBinding + ":" + serverPort, nil)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/fiam/gounidecode/unidecode"
	"github.com/fzzy/radix/extra/pool"
	"github.com/fzzy/radix/redis"
	"github.com/gorilla/mux"
	"github.com/mholt/binding"
	"github.com/namsral/flag"
	"github.com/unrolled/render"
	"net/http"
	"regexp"
	"runtime"
	"strings"
)

var junkRegex = regexp.MustCompile(`[\[\](){}"?!,-:;,']`)

func cleanWords(query string) []string {
	query = junkRegex.ReplaceAllString(query, " ")
	split := strings.Fields(strings.Trim(query, " "))
	terms := make([]string, len(split))
	var asciiTerm string
	for i, term := range split {
		terms[i] = strings.ToLower(strings.Trim(strings.Trim(term, " "), "."))
		asciiTerm = unidecode.Unidecode(terms[i])
		if asciiTerm != terms[i] {
			terms = append(terms, asciiTerm)
		}
	}
	return terms
}

func encodeString(str string) string {
	h := md5.New()
	h.Write([]byte(str))
	return hex.EncodeToString(h.Sum(nil))[0:8]
}

func getPrefixes(title string) []string {
	var prefixes []string
	for _, word := range cleanWords(title) {
		for i := 1; i <= len(word); i++ {
			prefixes = append(prefixes, word[0:i])
		}
	}
	return prefixes
}

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

func errHndlr(err error) {
	if err != nil {
		fmt.Println("error:", err)
		panic(err)
		// os.Exit(1)
	}
}

func IndexHandler(w http.ResponseWriter, req *http.Request) {
	// this assumes there's a `templates/index.tmpl` file
	renderer.HTML(w, http.StatusOK, "index", nil)
}

type UpdateForm struct {
	Domain     string
	Url        string
	Title      string
	Groups     string
	Popularity float64
}

func (f *UpdateForm) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&f.Domain: binding.Field{
			Form:     "domain",
			Required: true,
		},
		&f.Url: binding.Field{
			Form:     "url",
			Required: true,
		},
		&f.Title: binding.Field{
			Form:     "title",
			Required: true,
		},
		&f.Groups:     "groups",
		&f.Popularity: "popularity",
	}
}

func (f UpdateForm) Validate(req *http.Request, errs binding.Errors) binding.Errors {
	if strings.Trim(f.Domain, " ") == "" {
		errs = append(errs, binding.Error{
			FieldNames:     []string{"domain"},
			Classification: "ComplaintError",
			Message:        "Can't be empty",
		})
	}
	if strings.Trim(f.Title, " ") == "" {
		errs = append(errs, binding.Error{
			FieldNames:     []string{"title"},
			Classification: "ComplaintError",
			Message:        "Can't be empty",
		})
	}
	if strings.Trim(f.Url, " ") == "" {
		errs = append(errs, binding.Error{
			FieldNames:     []string{"url"},
			Classification: "ComplaintError",
			Message:        "Can't be empty",
		})
	}
	return errs
}

func UpdateHandler(w http.ResponseWriter, req *http.Request) {
	form := new(UpdateForm)
	errs := binding.Bind(req, form)
	if errs.Handle(w) {
		return
	}
	form.Domain = strings.Trim(form.Domain, " ")
	form.Title = strings.Trim(form.Title, " ")
	form.Url = strings.Trim(form.Url, " ")

	encoded := encodeString(form.Domain)
	url_encoded := encodeString(form.Url)

	c, err := redis_pool.Get()
	errHndlr(err)
	defer redis_pool.Put(c)
	// c.Cmd("FLUSHALL")
	// fmt.Println("CAREFUL! Always flushing the database")
	piped_commands := 0
	for _, prefix := range getPrefixes(form.Title) {
		c.Append("ZADD", encoded+prefix, form.Popularity, url_encoded)
		piped_commands += 1
	}
	c.Append("HSET", encoded+"$titles", url_encoded, form.Title)
	piped_commands += 1
	c.Append("HSET", encoded+"$urls", url_encoded, form.Url)
	piped_commands += 1
	for i := 1; i <= piped_commands; i++ {
		if err := c.GetReply().Err; err != nil {
			errHndlr(err)
		}
	}

	output := map[string]string{"message": "OK"}
	renderer.JSON(w, http.StatusCreated, output)
}

type FetchForm struct {
	Number int
	Query  string
	Domain string
}

func (f *FetchForm) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&f.Number: "n",
		&f.Query: binding.Field{
			Form:     "q",
			Required: true,
		},
		&f.Domain: binding.Field{
			Form:     "domain",
			Required: true,
		},
	}
}

func FetchHandler(w http.ResponseWriter, req *http.Request) {
	form := new(FetchForm)
	errs := binding.Bind(req, form)
	if errs.Handle(w) {
		return
	}
	n := form.Number
	if n <= 0 {
		n = 10 // default
	}

	form.Domain = strings.Trim(form.Domain, " ")

	encoded := encodeString(form.Domain)
	// fmt.Println(domain, encoded)

	form.Query = strings.Trim(form.Query, " ")
	terms := cleanWords(form.Query)

	c, err := redis_pool.Get()
	errHndlr(err)
	defer redis_pool.CarefullyPut(c, &err)

	encoded_terms := make([]string, len(terms))
	for i, term := range terms {
		encoded_terms[i] = encoded + term
	}
	// NOTE! Maybe we don't need the ZINTERSTORE if there's only 1 command
	c.Append("ZINTERSTORE", "$tmp", len(terms), encoded_terms, "AGGREGATE", "max")
	c.Append("ZREVRANGE", "$tmp", 0, n-1, "WITHSCORES")

	c.GetReply() // the ZINTERSTORE
	replies, err := c.GetReply().List()
	// fmt.Println("replies", replies, len(replies))
	errHndlr(err)

	encoded_urls := make([]string, n+1)
	scores := make([]string, n+1)
	evens := 0
	for i, element := range replies {
		if i%2 == 0 {
			encoded_urls[evens] = element
			evens = evens + 1
		} else {
			scores[evens-1] = element
		}
	}
	encoded_urls = encoded_urls[:evens]
	scores = scores[:evens]

	var titles []string
	var urls []string
	if len(encoded_urls) == 0 {
	} else {
		titles, err = c.Cmd("HMGET", encoded+"$titles", encoded_urls).List()
		errHndlr(err)
		urls, err = c.Cmd("HMGET", encoded+"$urls", encoded_urls).List()
		errHndlr(err)
	}
	rows := make([]interface{}, len(titles))
	for i, title := range titles {
		row := make([]string, 2)
		row[0] = urls[i]
		// fmt.Println("scores", scores[i]+ 1000)
		// row[1] = scores[i] * QueryScore(terms, title)
		row[1] = title
		rows[i] = row
	}
	rows = rows[:len(titles)]

	output := make(map[string]interface{})
	output["terms"] = terms
	output["results"] = rows
	// fmt.Println(output)
	w.Header().Set("Access-Control-Allow-Origin", "*")
	renderer.JSON(w, http.StatusOK, output)
}

var (
	renderer   *render.Render
	redis_pool *pool.Pool
	procs      int
	debug      = true
)

func main() {
	var port = 3001
	var redis_url = "127.0.0.1:6379"
	var redis_database = 0
	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(
		&redis_url, "redis_url", redis_url,
		"Redis URL to tcp connect to")
	flag.IntVar(&redis_database, "redis_database", redis_database,
		"Redis database number to connect to")
	flag.Parse()

	// Figuring out how many processors to use.
	max_procs := runtime.NumCPU()
	if procs == 0 {
		procs = max_procs
	} else if procs < 0 {
		panic("PROCS < 0")
	} else if procs > max_procs {
		panic(fmt.Sprintf("PROCS > max (%v)", max_procs))
	}
	// fmt.Println("procs=", procs)
	fmt.Printf("Running on %d procs\n", procs)
	runtime.GOMAXPROCS(procs)

	fmt.Println("DEBUG MODE:", debug)
	renderer = render.New(render.Options{
		IndentJSON:    debug,
		IsDevelopment: debug,
	})

	df := func(network, addr string) (*redis.Client, error) {
		client, err := redis.Dial(network, addr)
		// fmt.Println("DIaling")
		if err != nil {
			return nil, err
		}
		err = client.Cmd("SELECT", redis_database).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
	// fmt.Println("redis_url:", redis_url)
	// fmt.Println("redis_database:", redis_database)
	// fmt.Println("pool size", procs*10)
	redis_pool, err = pool.NewCustomPool("tcp", redis_url, 100, df)
	errHndlr(err)

	mux := mux.NewRouter()
	mux.HandleFunc("/", IndexHandler).Methods("GET", "HEAD")
	mux.HandleFunc("/v1", FetchHandler).Methods("GET", "HEAD")
	mux.HandleFunc("/v1", UpdateHandler).Methods("POST", "PUT")

	n := negroni.Classic()
	n.UseHandler(mux)
	n.Run(fmt.Sprintf(":%d", port))
}
<commit_msg>using a regex instead on cleanWords()<commit_after>package main

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/fiam/gounidecode/unidecode"
	"github.com/fzzy/radix/extra/pool"
	"github.com/fzzy/radix/redis"
	"github.com/gorilla/mux"
	"github.com/mholt/binding"
	"github.com/namsral/flag"
	"github.com/unrolled/render"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strings"
)

var junkRegex = regexp.MustCompile(`[\[\](){}"?!,-:;,']`)

func cleanWords(query string) []string {
	query = junkRegex.ReplaceAllString(query, " ")
	split := strings.Fields(strings.Trim(query, " "))
	terms := make([]string, len(split))
	var asciiTerm string
	for i, term := range split {
		terms[i] = strings.ToLower(strings.Trim(strings.Trim(term, " "), "."))
		asciiTerm = unidecode.Unidecode(terms[i])
		if asciiTerm != terms[i] {
			terms = append(terms, asciiTerm)
		}
	}
	return terms
}

func encodeString(str string) string {
	h := md5.New()
	h.Write([]byte(str))
	return hex.EncodeToString(h.Sum(nil))[0:8]
}

func getPrefixes(title string) []string {
	var prefixes []string
	for _, word := range cleanWords(title) {
		for i := 1; i <= len(word); i++ {
			prefixes = append(prefixes, word[0:i])
		}
	}
	return prefixes
}

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

func errHndlr(err error) {
	if err != nil {
		fmt.Println("error:", err)
		panic(err)
		os.Exit(1)
	}
}

func IndexHandler(w http.ResponseWriter, req *http.Request) {
	// this assumes there's a `templates/index.tmpl` file
	renderer.HTML(w, http.StatusOK, "index", nil)
}

type UpdateForm struct {
	Domain     string
	Url        string
	Title      string
	Groups     string
	Popularity float64
}

func (f *UpdateForm) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&f.Domain: binding.Field{
			Form:     "domain",
			Required: true,
		},
		&f.Url: binding.Field{
			Form:     "url",
			Required: true,
		},
		&f.Title: binding.Field{
			Form:     "title",
			Required: true,
		},
		&f.Groups:     "groups",
		&f.Popularity: "popularity",
	}
}

func (f UpdateForm) Validate(req *http.Request, errs binding.Errors) binding.Errors {
	if strings.Trim(f.Domain, " ") == "" {
		errs = append(errs, binding.Error{
			FieldNames:     []string{"domain"},
			Classification: "ComplaintError",
			Message:        "Can't be empty",
		})
	}
	if strings.Trim(f.Title, " ") == "" {
		errs = append(errs, binding.Error{
			FieldNames:     []string{"title"},
			Classification: "ComplaintError",
			Message:        "Can't be empty",
		})
	}
	if strings.Trim(f.Url, " ") == "" {
		errs = append(errs, binding.Error{
			FieldNames:     []string{"url"},
			Classification: "ComplaintError",
			Message:        "Can't be empty",
		})
	}
	return errs
}

func UpdateHandler(w http.ResponseWriter, req *http.Request) {
	form := new(UpdateForm)
	errs := binding.Bind(req, form)
	if errs.Handle(w) {
		return
	}
	form.Domain = strings.Trim(form.Domain, " ")
	form.Title = strings.Trim(form.Title, " ")
	form.Url = strings.Trim(form.Url, " ")

	encoded := encodeString(form.Domain)
	url_encoded := encodeString(form.Url)

	c, err := redis_pool.Get()
	errHndlr(err)
	defer redis_pool.Put(c)
	// c.Cmd("FLUSHALL")
	// fmt.Println("CAREFUL! Always flushing the database")
	piped_commands := 0
	for _, prefix := range getPrefixes(form.Title) {
		c.Append("ZADD", encoded+prefix, form.Popularity, url_encoded)
		piped_commands += 1
	}
	c.Append("HSET", encoded+"$titles", url_encoded, form.Title)
	piped_commands += 1
	c.Append("HSET", encoded+"$urls", url_encoded, form.Url)
	piped_commands += 1
	for i := 1; i <= piped_commands; i++ {
		if err := c.GetReply().Err; err != nil {
			errHndlr(err)
		}
	}

	output := map[string]string{"message": "OK"}
	renderer.JSON(w, http.StatusCreated, output)
}

type FetchForm struct {
	Number int
	Query  string
	Domain string
}

func (f *FetchForm) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&f.Number: "n",
		&f.Query: binding.Field{
			Form:     "q",
			Required: true,
		},
		&f.Domain: binding.Field{
			Form:     "domain",
			Required: true,
		},
	}
}

func FetchHandler(w http.ResponseWriter, req *http.Request) {
	form := new(FetchForm)
	errs := binding.Bind(req, form)
	if errs.Handle(w) {
		return
	}
	n := form.Number
	if n <= 0 {
		n = 10 // default
	}

	form.Domain = strings.Trim(form.Domain, " ")

	encoded := encodeString(form.Domain)
	// fmt.Println(domain, encoded)

	form.Query = strings.Trim(form.Query, " ")
	terms := cleanWords(form.Query)

	c, err := redis_pool.Get()
	errHndlr(err)
	defer redis_pool.CarefullyPut(c, &err)

	encoded_terms := make([]string, len(terms))
	for i, term := range terms {
		encoded_terms[i] = encoded + term
	}
	// NOTE! Maybe we don't need the ZINTERSTORE if there's only 1 command
	c.Append("ZINTERSTORE", "$tmp", len(terms), encoded_terms, "AGGREGATE", "max")
	c.Append("ZREVRANGE", "$tmp", 0, n-1, "WITHSCORES")

	c.GetReply() // the ZINTERSTORE
	replies, err := c.GetReply().List()
	// fmt.Println("replies", replies, len(replies))
	errHndlr(err)

	encoded_urls := make([]string, n+1)
	scores := make([]string, n+1)
	evens := 0
	for i, element := range replies {
		if i%2 == 0 {
			encoded_urls[evens] = element
			evens = evens + 1
		} else {
			scores[evens-1] = element
		}
	}
	encoded_urls = encoded_urls[:evens]
	scores = scores[:evens]

	var titles []string
	var urls []string
	if len(encoded_urls) == 0 {
	} else {
		titles, err = c.Cmd("HMGET", encoded+"$titles", encoded_urls).List()
		errHndlr(err)
		urls, err = c.Cmd("HMGET", encoded+"$urls", encoded_urls).List()
		errHndlr(err)
	}
	rows := make([]interface{}, len(titles))
	for i, title := range titles {
		row := make([]string, 2)
		row[0] = urls[i]
		// fmt.Println("scores", scores[i]+ 1000)
		// row[1] = scores[i] * QueryScore(terms, title)
		row[1] = title
		rows[i] = row
	}
	rows = rows[:len(titles)]

	output := make(map[string]interface{})
	output["terms"] = terms
	output["results"] = rows
	// fmt.Println(output)
	w.Header().Set("Access-Control-Allow-Origin", "*")
	renderer.JSON(w, http.StatusOK, output)
}

var (
	renderer   *render.Render
	redis_pool *pool.Pool
	procs      int
	debug      = true
)

func main() {
	var port = 3001
	var redis_url = "127.0.0.1:6379"
	var redis_database = 0
	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(
		&redis_url, "redis_url", redis_url,
		"Redis URL to tcp connect to")
	flag.IntVar(&redis_database, "redis_database", redis_database,
		"Redis database number to connect to")
	flag.Parse()

	// Figuring out how many processors to use.
	max_procs := runtime.NumCPU()
	if procs == 0 {
		procs = max_procs
	} else if procs < 0 {
		panic("PROCS < 0")
	} else if procs > max_procs {
		panic(fmt.Sprintf("PROCS > max (%v)", max_procs))
	}
	// fmt.Println("procs=", procs)
	fmt.Printf("Running on %d procs\n", procs)
	runtime.GOMAXPROCS(procs)

	fmt.Println("DEBUG MODE:", debug)
	renderer = render.New(render.Options{
		IndentJSON:    debug,
		IsDevelopment: debug,
	})

	df := func(network, addr string) (*redis.Client, error) {
		client, err := redis.Dial(network, addr)
		// fmt.Println("DIaling")
		if err != nil {
			return nil, err
		}
		err = client.Cmd("SELECT", redis_database).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
	// fmt.Println("redis_url:", redis_url)
	// fmt.Println("redis_database:", redis_database)
	// fmt.Println("pool size", procs*10)
	redis_pool, err = pool.NewCustomPool("tcp", redis_url, 100, df)
	errHndlr(err)

	mux := mux.NewRouter()
	mux.HandleFunc("/", IndexHandler).Methods("GET", "HEAD")
	mux.HandleFunc("/v1", FetchHandler).Methods("GET", "HEAD")
	mux.HandleFunc("/v1", UpdateHandler).Methods("POST", "PUT")

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


<commit_before>package sugoi

import (
	"log"
	"net/http"
)

func NewSugoi(port string) *SugoiServer {
	return &SugoiServer{
		handler: NewWrappedHandler(),
		port:    port,
	}
}

type SugoiServer struct {
	handler *WrappedHandler
	port    string
}

func (s *SugoiServer) SetStatic(url, dir string) {
	s.handler.staticUrl = url
	s.handler.staticDir = dir
}

func (s *SugoiServer) GetRoutes(method string) []*Route {
	routes := []*Route{}

	for _, o := range s.handler.routes {
		if o.method == method {
			routes = append(routes, o)
		}
	}

	return routes
}

func (s *SugoiServer) add(method string, path string, fn RouteHandler) *Route {
	route := CreateNewRoute(path, method, fn)
	s.handler.routes = append(s.handler.routes, route)

	return route
}

func (s *SugoiServer) Get(path string, fn RouteHandler) *Route {
	return s.add("get", path, fn)
}

func (s *SugoiServer) Delete(path string, fn RouteHandler) *Route {
	return s.add("delete", path, fn)
}

func (s *SugoiServer) Put(path string, fn RouteHandler) *Route {
	return s.add("put", path, fn)
}

func (s *SugoiServer) Post(path string, fn RouteHandler) *Route {
	return s.add("post", path, fn)
}

func (s *SugoiServer) Options(path string, fn RouteHandler) *Route {
	return s.add("options", path, fn)
}

func (s *SugoiServer) Patch(path string, fn RouteHandler) *Route {
	return s.add("patch", path, fn)
}

func Set404Page(s *SugoiServer, fn RouteHandler) {
	putDefaultHandler(s, http.StatusNotFound, fn)
}

func Set500Page(s *SugoiServer, fn RouteHandler) {
	putDefaultHandler(s, http.StatusInternalServerError, fn)
}

func Set401Page(s *SugoiServer, fn RouteHandler) {
	putDefaultHandler(s, http.StatusUnauthorized, fn)
}

func Set403Page(s *SugoiServer, fn RouteHandler) {
	putDefaultHandler(s, http.StatusForbidden, fn)
}

func putDefaultHandler(s *SugoiServer, code int, fn RouteHandler) {
	s.handler.defaultHandlers[code] = fn
}

func (s *SugoiServer) Serve() {
	err := http.ListenAndServe(":"+s.port, s.handler)
	if err != nil {
		log.Fatal(err)
	}
}

// Before every method
func (s *SugoiServer) Before(fn PreFilter) {
	s.handler.preFilters = append(s.handler.preFilters, fn)
}

// Handling Exceptions
func (s *SugoiServer) Error(fn ErrorHandler) {
	s.handler.errorHandlers = append(s.handler.errorHandlers, fn)
}
<commit_msg>add SugoiServer#Handle method<commit_after>package sugoi

import (
	"log"
	"net/http"
)

func NewSugoi(port string) *SugoiServer {
	return &SugoiServer{
		handler: NewWrappedHandler(),
		port:    port,
	}
}

type SugoiServer struct {
	handler *WrappedHandler
	port    string
}

func (s *SugoiServer) SetStatic(url, dir string) {
	s.handler.staticUrl = url
	s.handler.staticDir = dir
}

func (s *SugoiServer) GetRoutes(method string) []*Route {
	routes := []*Route{}

	for _, o := range s.handler.routes {
		if o.method == method {
			routes = append(routes, o)
		}
	}

	return routes
}

func (s *SugoiServer) add(method string, path string, fn RouteHandler) *Route {
	route := CreateNewRoute(path, method, fn)
	s.handler.routes = append(s.handler.routes, route)

	return route
}

func (s *SugoiServer) Handle(method string, path string, fn RouteHandler) *Route {
	return s.add(method, path, fn)
}

func (s *SugoiServer) Get(path string, fn RouteHandler) *Route {
	return s.Handle("get", path, fn)
}

func (s *SugoiServer) Delete(path string, fn RouteHandler) *Route {
	return s.Handle("delete", path, fn)
}

func (s *SugoiServer) Put(path string, fn RouteHandler) *Route {
	return s.Handle("put", path, fn)
}

func (s *SugoiServer) Post(path string, fn RouteHandler) *Route {
	return s.Handle("post", path, fn)
}

func (s *SugoiServer) Options(path string, fn RouteHandler) *Route {
	return s.Handle("options", path, fn)
}

func (s *SugoiServer) Patch(path string, fn RouteHandler) *Route {
	return s.Handle("patch", path, fn)
}

func Set404Page(s *SugoiServer, fn RouteHandler) {
	putDefaultHandler(s, http.StatusNotFound, fn)
}

func Set500Page(s *SugoiServer, fn RouteHandler) {
	putDefaultHandler(s, http.StatusInternalServerError, fn)
}

func Set401Page(s *SugoiServer, fn RouteHandler) {
	putDefaultHandler(s, http.StatusUnauthorized, fn)
}

func Set403Page(s *SugoiServer, fn RouteHandler) {
	putDefaultHandler(s, http.StatusForbidden, fn)
}

func putDefaultHandler(s *SugoiServer, code int, fn RouteHandler) {
	s.handler.defaultHandlers[code] = fn
}

func (s *SugoiServer) Serve() {
	err := http.ListenAndServe(":"+s.port, s.handler)
	if err != nil {
		log.Fatal(err)
	}
}

// Before every method
func (s *SugoiServer) Before(fn PreFilter) {
	s.handler.preFilters = append(s.handler.preFilters, fn)
}

// Handling Exceptions
func (s *SugoiServer) Error(fn ErrorHandler) {
	s.handler.errorHandlers = append(s.handler.errorHandlers, fn)
}
<|endoftext|>