<commit_before>// mgo - MongoDB driver for Go
//
// Copyright (c) 2010-2012 - Gustavo Niemeyer <gustavo@niemeyer.net>
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package mgo

import (
	"errors"
	"labix.org/v2/mgo/bson"
	"net"
	"sort"
	"sync"
	"time"
)

// ---------------------------------------------------------------------------
// Mongo server encapsulation.

type mongoServer struct {
	sync.RWMutex
	Addr          string
	ResolvedAddr  string
	tcpaddr       *net.TCPAddr
	unusedSockets []*mongoSocket
	liveSockets   []*mongoSocket
	closed        bool
	master        bool
	abended       bool
	sync          chan bool
	dial          dialer
	pingValue     time.Duration
	pingIndex     int
	pingCount     uint32
	pingWindow    [6]time.Duration
}

type dialer func(addr net.Addr) (net.Conn, error)

func newServer(addr string, tcpaddr *net.TCPAddr, sync chan bool, dial dialer) *mongoServer {
	server := &mongoServer{
		Addr:         addr,
		ResolvedAddr: tcpaddr.String(),
		tcpaddr:      tcpaddr,
		sync:         sync,
		dial:         dial,
	}
	// Once so the server gets a ping value, then loop in background.
	server.pinger(false)
	go server.pinger(true)
	return server
}

var errSocketLimit = errors.New("per-server connection limit reached")
var errServerClosed = errors.New("server was closed")

// AcquireSocket returns a socket for communicating with the server.
// This will attempt to reuse an old connection, if one is available. Otherwise,
// it will establish a new one. The returned socket is owned by the call site,
// and will return to the cache when the socket has its Release method called
// the same number of times as AcquireSocket + Acquire were called for it.
// If the limit argument is not zero, a socket will only be returned if the
// number of sockets in use for this server is under the provided limit.
func (server *mongoServer) AcquireSocket(limit int) (socket *mongoSocket, abended bool, err error) {
	for {
		server.Lock()
		abended = server.abended
		n := len(server.unusedSockets)
		if limit > 0 && len(server.liveSockets)-n >= limit {
			server.Unlock()
			return nil, false, errSocketLimit
		}
		if n > 0 {
			socket = server.unusedSockets[n-1]
			server.unusedSockets[n-1] = nil // Help GC.
			server.unusedSockets = server.unusedSockets[:n-1]
			server.Unlock()
			err = socket.InitialAcquire()
			if err != nil {
				continue
			}
		} else {
			server.Unlock()
			socket, err = server.Connect()
			if err == nil {
				server.Lock()
				if server.closed {
					socket.Close()
					return nil, false, errServerClosed
				}
				server.liveSockets = append(server.liveSockets, socket)
				server.Unlock()
			}
		}
		return
	}
	panic("unreachable")
}

// Connect establishes a new connection to the server. This should
// generally be done through server.AcquireSocket().
func (server *mongoServer) Connect() (*mongoSocket, error) {
	server.RLock()
	master := server.master
	dial := server.dial
	server.RUnlock()

	log("Establishing new connection to ", server.Addr, "...")
	var conn net.Conn
	var err error
	if dial == nil {
		conn, err = net.DialTCP("tcp", nil, server.tcpaddr)
	} else {
		conn, err = dial(server.tcpaddr)
	}
	if err != nil {
		log("Connection to ", server.Addr, " failed: ", err.Error())
		return nil, err
	}
	log("Connection to ", server.Addr, " established.")

	stats.conn(+1, master)
	return newSocket(server, conn), nil
}

// Close forces closing all sockets that are alive, whether
// they're currently in use or not.
func (server *mongoServer) Close() {
	server.Lock()
	server.closed = true
	liveSockets := server.liveSockets
	unusedSockets := server.unusedSockets
	server.liveSockets = nil
	server.unusedSockets = nil
	server.Unlock()
	logf("Connections to %s closing (%d live sockets).", server.Addr, len(liveSockets))
	for i, s := range liveSockets {
		s.Close()
		liveSockets[i] = nil
	}
	for i := range unusedSockets {
		unusedSockets[i] = nil
	}
}

// RecycleSocket puts socket back into the unused cache.
func (server *mongoServer) RecycleSocket(socket *mongoSocket) {
	server.Lock()
	if !server.closed {
		server.unusedSockets = append(server.unusedSockets, socket)
	}
	server.Unlock()
}

func removeSocket(sockets []*mongoSocket, socket *mongoSocket) []*mongoSocket {
	for i, s := range sockets {
		if s == socket {
			copy(sockets[i:], sockets[i+1:])
			n := len(sockets) - 1
			sockets[n] = nil
			sockets = sockets[:n]
			break
		}
	}
	return sockets
}

// AbendSocket notifies the server that the given socket has terminated
// abnormally, and thus should be discarded rather than cached.
func (server *mongoServer) AbendSocket(socket *mongoSocket) {
	server.Lock()
	server.abended = true
	if server.closed {
		server.Unlock()
		return
	}
	server.liveSockets = removeSocket(server.liveSockets, socket)
	server.unusedSockets = removeSocket(server.unusedSockets, socket)
	server.Unlock()
	// Maybe just a timeout, but suggest a cluster sync up just in case.
	select {
	case server.sync <- true:
	default:
	}
}

func (server *mongoServer) SetMaster(isMaster bool) {
	server.Lock()
	server.master = isMaster
	server.Unlock()
}

func (server *mongoServer) IsMaster() bool {
	server.RLock()
	result := server.master
	server.RUnlock()
	return result
}

var pingDelay = 10 * time.Second

func (server *mongoServer) pinger(loop bool) {
	op := queryOp{
		collection: "admin.$cmd",
		query:      bson.D{{"ping", 1}},
		flags:      flagSlaveOk,
		limit:      -1,
	}
	for {
		if loop {
			time.Sleep(pingDelay)
			server.RLock()
			closed := server.closed
			server.RUnlock()
			if closed {
				return
			}
		}
		op := op
		socket, _, err := server.AcquireSocket(0)
		if err == nil {
			start := time.Now()
			_, _ = socket.SimpleQuery(&op)
			delay := time.Now().Sub(start)

			server.pingWindow[server.pingIndex] = delay
			server.pingIndex = (server.pingIndex + 1) % len(server.pingWindow)
			server.pingCount++
			var max time.Duration
			for i := 0; i < len(server.pingWindow) && uint32(i) < server.pingCount; i++ {
				if server.pingWindow[i] > max {
					max = server.pingWindow[i]
				}
			}
			socket.Release()
			server.Lock()
			if server.closed {
				loop = false
			}
			server.pingValue = max
			server.Unlock()
			logf("Ping for %s is %d ms", server.Addr, max/time.Millisecond)
		} else if err == errServerClosed {
			return
		}
		if !loop {
			return
		}
	}
}

type mongoServerSlice []*mongoServer

func (s mongoServerSlice) Len() int {
	return len(s)
}

func (s mongoServerSlice) Less(i, j int) bool {
	return s[i].ResolvedAddr < s[j].ResolvedAddr
}

func (s mongoServerSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s mongoServerSlice) Sort() {
	sort.Sort(s)
}

func (s mongoServerSlice) Search(resolvedAddr string) (i int, ok bool) {
	n := len(s)
	i = sort.Search(n, func(i int) bool {
		return s[i].ResolvedAddr >= resolvedAddr
	})
	return i, i != n && s[i].ResolvedAddr == resolvedAddr
}

type mongoServers struct {
	slice mongoServerSlice
}

func (servers *mongoServers) Search(resolvedAddr string) (server *mongoServer) {
	if i, ok := servers.slice.Search(resolvedAddr); ok {
		return servers.slice[i]
	}
	return nil
}

func (servers *mongoServers) Add(server *mongoServer) {
	servers.slice = append(servers.slice, server)
	servers.slice.Sort()
}

func (servers *mongoServers) Remove(other *mongoServer) (server *mongoServer) {
	if i, found := servers.slice.Search(other.ResolvedAddr); found {
		server = servers.slice[i]
		copy(servers.slice[i:], servers.slice[i+1:])
		n := len(servers.slice) - 1
		servers.slice[n] = nil // Help GC.
		servers.slice = servers.slice[:n]
	}
	return
}

func (servers *mongoServers) Slice() []*mongoServer {
	return ([]*mongoServer)(servers.slice)
}

func (servers *mongoServers) Get(i int) *mongoServer {
	return servers.slice[i]
}

func (servers *mongoServers) Len() int {
	return len(servers.slice)
}

func (servers *mongoServers) Empty() bool {
	return len(servers.slice) == 0
}

// MostAvailable returns the best guess of what would be the
// most interesting server to perform operations on at this
// point in time.
func (servers *mongoServers) MostAvailable() *mongoServer {
	if len(servers.slice) == 0 {
		panic("MostAvailable: can't be used on empty server list")
	}
	var best *mongoServer
	for i, next := range servers.slice {
		if i == 0 {
			best = next
			best.RLock()
			continue
		}
		next.RLock()
		swap := false
		switch {
		case next.master != best.master:
			// Prefer slaves.
			swap = best.master
		case next.pingValue < best.pingValue-15*time.Millisecond:
			// Prefer nearest server.
			swap = true
		case len(next.liveSockets)-len(next.unusedSockets) < len(best.liveSockets)-len(best.unusedSockets):
			// Prefer servers with less connections.
			swap = true
		}
		if swap {
			best.RUnlock()
			best = next
		} else {
			next.RUnlock()
		}
	}
	best.RUnlock()
	return best
}
<commit_msg>Always unlock the server<commit_after>// mgo - MongoDB driver for Go
//
// Copyright (c) 2010-2012 - Gustavo Niemeyer <gustavo@niemeyer.net>
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package mgo

import (
	"errors"
	"labix.org/v2/mgo/bson"
	"net"
	"sort"
	"sync"
	"time"
)

// ---------------------------------------------------------------------------
// Mongo server encapsulation.

type mongoServer struct {
	sync.RWMutex
	Addr          string
	ResolvedAddr  string
	tcpaddr       *net.TCPAddr
	unusedSockets []*mongoSocket
	liveSockets   []*mongoSocket
	closed        bool
	master        bool
	abended       bool
	sync          chan bool
	dial          dialer
	pingValue     time.Duration
	pingIndex     int
	pingCount     uint32
	pingWindow    [6]time.Duration
}

type dialer func(addr net.Addr) (net.Conn, error)

func newServer(addr string, tcpaddr *net.TCPAddr, sync chan bool, dial dialer) *mongoServer {
	server := &mongoServer{
		Addr:         addr,
		ResolvedAddr: tcpaddr.String(),
		tcpaddr:      tcpaddr,
		sync:         sync,
		dial:         dial,
	}
	// Once so the server gets a ping value, then loop in background.
	server.pinger(false)
	go server.pinger(true)
	return server
}

var errSocketLimit = errors.New("per-server connection limit reached")
var errServerClosed = errors.New("server was closed")

// AcquireSocket returns a socket for communicating with the server.
// This will attempt to reuse an old connection, if one is available. Otherwise,
// it will establish a new one. The returned socket is owned by the call site,
// and will return to the cache when the socket has its Release method called
// the same number of times as AcquireSocket + Acquire were called for it.
// If the limit argument is not zero, a socket will only be returned if the
// number of sockets in use for this server is under the provided limit.
func (server *mongoServer) AcquireSocket(limit int) (socket *mongoSocket, abended bool, err error) {
	for {
		server.Lock()
		abended = server.abended
		n := len(server.unusedSockets)
		if limit > 0 && len(server.liveSockets)-n >= limit {
			server.Unlock()
			return nil, false, errSocketLimit
		}
		if n > 0 {
			socket = server.unusedSockets[n-1]
			server.unusedSockets[n-1] = nil // Help GC.
			server.unusedSockets = server.unusedSockets[:n-1]
			server.Unlock()
			err = socket.InitialAcquire()
			if err != nil {
				continue
			}
		} else {
			server.Unlock()
			socket, err = server.Connect()
			if err == nil {
				server.Lock()
				if server.closed {
					socket.Close()
					socket = nil
					err = errServerClosed
				} else {
					server.liveSockets = append(server.liveSockets, socket)
				}
				server.Unlock()
			}
		}
		return
	}
	panic("unreachable")
}

// Connect establishes a new connection to the server. This should
// generally be done through server.AcquireSocket().
func (server *mongoServer) Connect() (*mongoSocket, error) {
	server.RLock()
	master := server.master
	dial := server.dial
	server.RUnlock()

	log("Establishing new connection to ", server.Addr, "...")
	var conn net.Conn
	var err error
	if dial == nil {
		conn, err = net.DialTCP("tcp", nil, server.tcpaddr)
	} else {
		conn, err = dial(server.tcpaddr)
	}
	if err != nil {
		log("Connection to ", server.Addr, " failed: ", err.Error())
		return nil, err
	}
	log("Connection to ", server.Addr, " established.")

	stats.conn(+1, master)
	return newSocket(server, conn), nil
}

// Close forces closing all sockets that are alive, whether
// they're currently in use or not.
func (server *mongoServer) Close() {
	server.Lock()
	server.closed = true
	liveSockets := server.liveSockets
	unusedSockets := server.unusedSockets
	server.liveSockets = nil
	server.unusedSockets = nil
	server.Unlock()
	logf("Connections to %s closing (%d live sockets).", server.Addr, len(liveSockets))
	for i, s := range liveSockets {
		s.Close()
		liveSockets[i] = nil
	}
	for i := range unusedSockets {
		unusedSockets[i] = nil
	}
}

// RecycleSocket puts socket back into the unused cache.
func (server *mongoServer) RecycleSocket(socket *mongoSocket) {
	server.Lock()
	if !server.closed {
		server.unusedSockets = append(server.unusedSockets, socket)
	}
	server.Unlock()
}

func removeSocket(sockets []*mongoSocket, socket *mongoSocket) []*mongoSocket {
	for i, s := range sockets {
		if s == socket {
			copy(sockets[i:], sockets[i+1:])
			n := len(sockets) - 1
			sockets[n] = nil
			sockets = sockets[:n]
			break
		}
	}
	return sockets
}

// AbendSocket notifies the server that the given socket has terminated
// abnormally, and thus should be discarded rather than cached.
func (server *mongoServer) AbendSocket(socket *mongoSocket) {
	server.Lock()
	server.abended = true
	if server.closed {
		server.Unlock()
		return
	}
	server.liveSockets = removeSocket(server.liveSockets, socket)
	server.unusedSockets = removeSocket(server.unusedSockets, socket)
	server.Unlock()
	// Maybe just a timeout, but suggest a cluster sync up just in case.
	select {
	case server.sync <- true:
	default:
	}
}

func (server *mongoServer) SetMaster(isMaster bool) {
	server.Lock()
	server.master = isMaster
	server.Unlock()
}

func (server *mongoServer) IsMaster() bool {
	server.RLock()
	result := server.master
	server.RUnlock()
	return result
}

var pingDelay = 10 * time.Second

func (server *mongoServer) pinger(loop bool) {
	op := queryOp{
		collection: "admin.$cmd",
		query:      bson.D{{"ping", 1}},
		flags:      flagSlaveOk,
		limit:      -1,
	}
	for {
		if loop {
			time.Sleep(pingDelay)
			server.RLock()
			closed := server.closed
			server.RUnlock()
			if closed {
				return
			}
		}
		op := op
		socket, _, err := server.AcquireSocket(0)
		if err == nil {
			start := time.Now()
			_, _ = socket.SimpleQuery(&op)
			delay := time.Now().Sub(start)

			server.pingWindow[server.pingIndex] = delay
			server.pingIndex = (server.pingIndex + 1) % len(server.pingWindow)
			server.pingCount++
			var max time.Duration
			for i := 0; i < len(server.pingWindow) && uint32(i) < server.pingCount; i++ {
				if server.pingWindow[i] > max {
					max = server.pingWindow[i]
				}
			}
			socket.Release()
			server.Lock()
			if server.closed {
				loop = false
			}
			server.pingValue = max
			server.Unlock()
			logf("Ping for %s is %d ms", server.Addr, max/time.Millisecond)
		} else if err == errServerClosed {
			return
		}
		if !loop {
			return
		}
	}
}

type mongoServerSlice []*mongoServer

func (s mongoServerSlice) Len() int {
	return len(s)
}

func (s mongoServerSlice) Less(i, j int) bool {
	return s[i].ResolvedAddr < s[j].ResolvedAddr
}

func (s mongoServerSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s mongoServerSlice) Sort() {
	sort.Sort(s)
}

func (s mongoServerSlice) Search(resolvedAddr string) (i int, ok bool) {
	n := len(s)
	i = sort.Search(n, func(i int) bool {
		return s[i].ResolvedAddr >= resolvedAddr
	})
	return i, i != n && s[i].ResolvedAddr == resolvedAddr
}

type mongoServers struct {
	slice mongoServerSlice
}

func (servers *mongoServers) Search(resolvedAddr string) (server *mongoServer) {
	if i, ok := servers.slice.Search(resolvedAddr); ok {
		return servers.slice[i]
	}
	return nil
}

func (servers *mongoServers) Add(server *mongoServer) {
	servers.slice = append(servers.slice, server)
	servers.slice.Sort()
}

func (servers *mongoServers) Remove(other *mongoServer) (server *mongoServer) {
	if i, found := servers.slice.Search(other.ResolvedAddr); found {
		server = servers.slice[i]
		copy(servers.slice[i:], servers.slice[i+1:])
		n := len(servers.slice) - 1
		servers.slice[n] = nil // Help GC.
		servers.slice = servers.slice[:n]
	}
	return
}

func (servers *mongoServers) Slice() []*mongoServer {
	return ([]*mongoServer)(servers.slice)
}

func (servers *mongoServers) Get(i int) *mongoServer {
	return servers.slice[i]
}

func (servers *mongoServers) Len() int {
	return len(servers.slice)
}

func (servers *mongoServers) Empty() bool {
	return len(servers.slice) == 0
}

// MostAvailable returns the best guess of what would be the
// most interesting server to perform operations on at this
// point in time.
func (servers *mongoServers) MostAvailable() *mongoServer {
	if len(servers.slice) == 0 {
		panic("MostAvailable: can't be used on empty server list")
	}
	var best *mongoServer
	for i, next := range servers.slice {
		if i == 0 {
			best = next
			best.RLock()
			continue
		}
		next.RLock()
		swap := false
		switch {
		case next.master != best.master:
			// Prefer slaves.
			swap = best.master
		case next.pingValue < best.pingValue-15*time.Millisecond:
			// Prefer nearest server.
			swap = true
		case len(next.liveSockets)-len(next.unusedSockets) < len(best.liveSockets)-len(best.unusedSockets):
			// Prefer servers with less connections.
			swap = true
		}
		if swap {
			best.RUnlock()
			best = next
		} else {
			next.RUnlock()
		}
	}
	best.RUnlock()
	return best
}
<|endoftext|>


<commit_before>// Copyright 2013 Gary Burd. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package websocket

import (
	"bufio"
	"errors"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// HandshakeError describes an error with the handshake from the peer.
type HandshakeError struct {
	message string
}

func (e HandshakeError) Error() string { return e.message }

const (
	defaultReadBufferSize  = 4096
	defaultWriteBufferSize = 4096
)

type Upgrader struct {
	// HandshakeTimeout specifies the duration for the handshake to complete.
	HandshakeTimeout time.Duration

	// Input and output buffer sizes. If the buffer size is zero, then
	// default values will be used.
	ReadBufferSize, WriteBufferSize int

	// Subprotocols specifies the server's supported protocols. If Subprotocols
	// is nil, then Upgrade does not negotiate a subprotocol.
	Subprotocols []string

	// Error specifies the function for generating HTTP error responses. If Error
	// is nil, then http.Error is used to generate the HTTP response.
	Error func(w http.ResponseWriter, r *http.Request, status int, reason error)

	// CheckOrigin returns true if the request Origin header is acceptable.
	// If CheckOrigin is nil, the host in the Origin header must match
	// the host of the request.
	CheckOrigin func(r *http.Request) bool
}

// Return an error depending on settings on the Upgrader
func (u *Upgrader) returnError(w http.ResponseWriter, r *http.Request, status int, reason error) {
	if u.Error != nil {
		u.Error(w, r, status, reason)
	} else {
		http.Error(w, reason.Error(), status)
	}
}

// Check if the passed subprotocol is supported by the server
func (u *Upgrader) hasSubprotocol(subprotocol string) bool {
	if u.Subprotocols == nil {
		return false
	}

	for _, s := range u.Subprotocols {
		if s == subprotocol {
			return true
		}
	}

	return false
}

// Check if host in Origin header matches host of request
func (u *Upgrader) checkSameOrigin(r *http.Request) bool {
	origin := r.Header.Get("Origin")
	if origin == "" {
		return false
	}
	uri, err := url.ParseRequestURI(origin)
	if err != nil {
		return false
	}
	return uri.Host == r.Host
}

// Upgrade upgrades the HTTP server connection to the WebSocket protocol.
//
// The responseHeader is included in the response to the client's upgrade
// request. Use the responseHeader to specify cookies (Set-Cookie).
//
// The connection buffers IO to the underlying network connection.
// Messages can be larger than the buffers.
//
// If the request is not a valid WebSocket handshake, then Upgrade returns an
// error of type HandshakeError. Depending on settings on the Upgrader,
// an error message already has been returned to the caller.
func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*Conn, error) {
	if values := r.Header["Sec-Websocket-Version"]; len(values) == 0 || values[0] != "13" {
		err := HandshakeError{"websocket: version != 13"}
		u.returnError(w, r, http.StatusBadRequest, err)
		return nil, err
	}

	if !tokenListContainsValue(r.Header, "Connection", "upgrade") {
		err := HandshakeError{"websocket: connection header != upgrade"}
		u.returnError(w, r, http.StatusBadRequest, err)
		return nil, err
	}

	if !tokenListContainsValue(r.Header, "Upgrade", "websocket") {
		err := HandshakeError{"websocket: upgrade != websocket"}
		u.returnError(w, r, http.StatusBadRequest, err)
		return nil, err
	}

	checkOrigin := u.CheckOrigin
	if checkOrigin == nil {
		checkOrigin = u.checkSameOrigin
	}
	if !checkOrigin(r) {
		err := HandshakeError{"websocket: origin not allowed"}
		u.returnError(w, r, http.StatusForbidden, err)
		return nil, err
	}

	var challengeKey string
	values := r.Header["Sec-Websocket-Key"]
	if len(values) == 0 || values[0] == "" {
		err := HandshakeError{"websocket: key missing or blank"}
		u.returnError(w, r, http.StatusBadRequest, err)
		return nil, err
	}
	challengeKey = values[0]

	var (
		netConn net.Conn
		br      *bufio.Reader
		err     error
	)

	h, ok := w.(http.Hijacker)
	if !ok {
		return nil, errors.New("websocket: response does not implement http.Hijacker")
	}
	var rw *bufio.ReadWriter
	netConn, rw, err = h.Hijack()
	br = rw.Reader

	if br.Buffered() > 0 {
		netConn.Close()
		return nil, errors.New("websocket: client sent data before handshake is complete")
	}

	readBufSize := u.ReadBufferSize
	if readBufSize == 0 {
		readBufSize = defaultReadBufferSize
	}
	writeBufSize := u.WriteBufferSize
	if writeBufSize == 0 {
		writeBufSize = defaultWriteBufferSize
	}
	c := newConn(netConn, true, readBufSize, writeBufSize)

	if u.Subprotocols != nil {
		for _, proto := range Subprotocols(r) {
			if u.hasSubprotocol(proto) {
				c.subprotocol = proto
				break
			}
		}
	} else if responseHeader != nil {
		c.subprotocol = responseHeader.Get("Sec-Websocket-Protocol")
	}

	p := c.writeBuf[:0]
	p = append(p, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: "...)
	p = append(p, computeAcceptKey(challengeKey)...)
	p = append(p, "\r\n"...)
	if c.subprotocol != "" {
		p = append(p, "Sec-Websocket-Protocol: "...)
		p = append(p, c.subprotocol...)
		p = append(p, "\r\n"...)
	}
	for k, vs := range responseHeader {
		if k == "Sec-Websocket-Protocol" {
			continue
		}
		for _, v := range vs {
			p = append(p, k...)
			p = append(p, ": "...)
			for i := 0; i < len(v); i++ {
				b := v[i]
				if b <= 31 {
					// prevent response splitting.
					b = ' '
				}
				p = append(p, b)
			}
			p = append(p, "\r\n"...)
		}
	}
	p = append(p, "\r\n"...)

	if u.HandshakeTimeout > 0 {
		netConn.SetWriteDeadline(time.Now().Add(u.HandshakeTimeout))
	}
	if _, err = netConn.Write(p); err != nil {
		netConn.Close()
		return nil, err
	}

	return c, nil
}

// This method is deprecated, use websocket.Upgrader instead.
//
// Upgrade upgrades the HTTP server connection to the WebSocket protocol.
//
// The application is responsible for checking the request origin before
// calling Upgrade. An example implementation of the same origin policy is:
//
//	if req.Header.Get("Origin") != "http://"+req.Host {
//		http.Error(w, "Origin not allowed", 403)
//		return
//	}
//
// If the endpoint supports subprotocols, then the application is responsible
// for negotiating the protocol used on the connection. Use the Subprotocols()
// function to get the subprotocols requested by the client. Use the
// Sec-Websocket-Protocol response header to specify the subprotocol selected
// by the application.
//
// The responseHeader is included in the response to the client's upgrade
// request. Use the responseHeader to specify cookies (Set-Cookie) and the
// negotiated subprotocol (Sec-Websocket-Protocol).
//
// The connection buffers IO to the underlying network connection. The
// readBufSize and writeBufSize parameters specify the size of the buffers to
// use. Messages can be larger than the buffers.
//
// If the request is not a valid WebSocket handshake, then Upgrade returns an
// error of type HandshakeError. Applications should handle this error by
// replying to the client with an HTTP error response.
func Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header, readBufSize, writeBufSize int) (*Conn, error) {
	u := Upgrader{ReadBufferSize: readBufSize, WriteBufferSize: writeBufSize}
	u.Error = func(w http.ResponseWriter, r *http.Request, status int, reason error) {
		// don't return errors to maintain backwards compatibility
	}
	u.CheckOrigin = func(r *http.Request) bool {
		// allow all connections by default
		return true
	}
	return u.Upgrade(w, r, responseHeader)
}

// Subprotocols returns the subprotocols requested by the client in the
// Sec-Websocket-Protocol header.
func Subprotocols(r *http.Request) []string {
	h := strings.TrimSpace(r.Header.Get("Sec-Websocket-Protocol"))
	if h == "" {
		return nil
	}
	protocols := strings.Split(h, ",")
	for i := range protocols {
		protocols[i] = strings.TrimSpace(protocols[i])
	}
	return protocols
}
<commit_msg>Simplified returning of errors.<commit_after>// Copyright 2013 Gary Burd. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package websocket

import (
	"bufio"
	"errors"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// HandshakeError describes an error with the handshake from the peer.
type HandshakeError struct {
	message string
}

func (e HandshakeError) Error() string { return e.message }

const (
	defaultReadBufferSize  = 4096
	defaultWriteBufferSize = 4096
)

type Upgrader struct {
	// HandshakeTimeout specifies the duration for the handshake to complete.
	HandshakeTimeout time.Duration

	// Input and output buffer sizes. If the buffer size is zero, then
	// default values will be used.
	ReadBufferSize, WriteBufferSize int

	// Subprotocols specifies the server's supported protocols. If Subprotocols
	// is nil, then Upgrade does not negotiate a subprotocol.
	Subprotocols []string

	// Error specifies the function for generating HTTP error responses. If Error
	// is nil, then http.Error is used to generate the HTTP response.
	Error func(w http.ResponseWriter, r *http.Request, status int, reason error)

	// CheckOrigin returns true if the request Origin header is acceptable.
	// If CheckOrigin is nil, the host in the Origin header must match
	// the host of the request.
	CheckOrigin func(r *http.Request) bool
}

// Return an error depending on settings on the Upgrader
func (u *Upgrader) returnError(w http.ResponseWriter, r *http.Request, status int, reason string) (*Conn, error) {
	err := HandshakeError{reason}
	if u.Error != nil {
		u.Error(w, r, status, err)
	} else {
		http.Error(w, reason, status)
	}
	return nil, err
}

// Check if the passed subprotocol is supported by the server
func (u *Upgrader) hasSubprotocol(subprotocol string) bool {
	if u.Subprotocols == nil {
		return false
	}

	for _, s := range u.Subprotocols {
		if s == subprotocol {
			return true
		}
	}

	return false
}

// Check if host in Origin header matches host of request
func (u *Upgrader) checkSameOrigin(r *http.Request) bool {
	origin := r.Header.Get("Origin")
	if origin == "" {
		return false
	}
	uri, err := url.ParseRequestURI(origin)
	if err != nil {
		return false
	}
	return uri.Host == r.Host
}

// Upgrade upgrades the HTTP server connection to the WebSocket protocol.
//
// The responseHeader is included in the response to the client's upgrade
// request. Use the responseHeader to specify cookies (Set-Cookie).
//
// The connection buffers IO to the underlying network connection.
// Messages can be larger than the buffers.
//
// If the request is not a valid WebSocket handshake, then Upgrade returns an
// error of type HandshakeError. Depending on settings on the Upgrader,
// an error message already has been returned to the caller.
func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*Conn, error) {
	if values := r.Header["Sec-Websocket-Version"]; len(values) == 0 || values[0] != "13" {
		return u.returnError(w, r, http.StatusBadRequest, "websocket: version != 13")
	}

	if !tokenListContainsValue(r.Header, "Connection", "upgrade") {
		return u.returnError(w, r, http.StatusBadRequest, "websocket: connection header != upgrade")
	}

	if !tokenListContainsValue(r.Header, "Upgrade", "websocket") {
		return u.returnError(w, r, http.StatusBadRequest, "websocket: upgrade != websocket")
	}

	checkOrigin := u.CheckOrigin
	if checkOrigin == nil {
		checkOrigin = u.checkSameOrigin
	}
	if !checkOrigin(r) {
		return u.returnError(w, r, http.StatusForbidden, "websocket: origin not allowed")
	}

	var challengeKey string
	values := r.Header["Sec-Websocket-Key"]
	if len(values) == 0 || values[0] == "" {
		return u.returnError(w, r, http.StatusBadRequest, "websocket: key missing or blank")
	}
	challengeKey = values[0]

	var (
		netConn net.Conn
		br      *bufio.Reader
		err     error
	)

	h, ok := w.(http.Hijacker)
	if !ok {
		return nil, errors.New("websocket: response does not implement http.Hijacker")
	}
	var rw *bufio.ReadWriter
	netConn, rw, err = h.Hijack()
	br = rw.Reader

	if br.Buffered() > 0 {
		netConn.Close()
		return nil, errors.New("websocket: client sent data before handshake is complete")
	}

	readBufSize := u.ReadBufferSize
	if readBufSize == 0 {
		readBufSize = defaultReadBufferSize
	}
	writeBufSize := u.WriteBufferSize
	if writeBufSize == 0 {
		writeBufSize = defaultWriteBufferSize
	}
	c := newConn(netConn, true, readBufSize, writeBufSize)

	if u.Subprotocols != nil {
		for _, proto := range Subprotocols(r) {
			if u.hasSubprotocol(proto) {
				c.subprotocol = proto
				break
			}
		}
	} else if responseHeader != nil {
		c.subprotocol = responseHeader.Get("Sec-Websocket-Protocol")
	}

	p := c.writeBuf[:0]
	p = append(p, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: "...)
	p = append(p, computeAcceptKey(challengeKey)...)
	p = append(p, "\r\n"...)
	if c.subprotocol != "" {
		p = append(p, "Sec-Websocket-Protocol: "...)
		p = append(p, c.subprotocol...)
		p = append(p, "\r\n"...)
	}
	for k, vs := range responseHeader {
		if k == "Sec-Websocket-Protocol" {
			continue
		}
		for _, v := range vs {
			p = append(p, k...)
			p = append(p, ": "...)
			for i := 0; i < len(v); i++ {
				b := v[i]
				if b <= 31 {
					// prevent response splitting.
					b = ' '
				}
				p = append(p, b)
			}
			p = append(p, "\r\n"...)
		}
	}
	p = append(p, "\r\n"...)

	if u.HandshakeTimeout > 0 {
		netConn.SetWriteDeadline(time.Now().Add(u.HandshakeTimeout))
	}
	if _, err = netConn.Write(p); err != nil {
		netConn.Close()
		return nil, err
	}

	return c, nil
}

// This method is deprecated, use websocket.Upgrader instead.
//
// Upgrade upgrades the HTTP server connection to the WebSocket protocol.
//
// The application is responsible for checking the request origin before
// calling Upgrade. An example implementation of the same origin policy is:
//
//	if req.Header.Get("Origin") != "http://"+req.Host {
//		http.Error(w, "Origin not allowed", 403)
//		return
//	}
//
// If the endpoint supports subprotocols, then the application is responsible
// for negotiating the protocol used on the connection. Use the Subprotocols()
// function to get the subprotocols requested by the client. Use the
// Sec-Websocket-Protocol response header to specify the subprotocol selected
// by the application.
//
// The responseHeader is included in the response to the client's upgrade
// request. Use the responseHeader to specify cookies (Set-Cookie) and the
// negotiated subprotocol (Sec-Websocket-Protocol).
//
// The connection buffers IO to the underlying network connection. The
// readBufSize and writeBufSize parameters specify the size of the buffers to
// use. Messages can be larger than the buffers.
//
// If the request is not a valid WebSocket handshake, then Upgrade returns an
// error of type HandshakeError. Applications should handle this error by
// replying to the client with an HTTP error response.
func Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header, readBufSize, writeBufSize int) (*Conn, error) {
	u := Upgrader{ReadBufferSize: readBufSize, WriteBufferSize: writeBufSize}
	u.Error = func(w http.ResponseWriter, r *http.Request, status int, reason error) {
		// don't return errors to maintain backwards compatibility
	}
	u.CheckOrigin = func(r *http.Request) bool {
		// allow all connections by default
		return true
	}
	return u.Upgrade(w, r, responseHeader)
}

// Subprotocols returns the subprotocols requested by the client in the
// Sec-Websocket-Protocol header.
func Subprotocols(r *http.Request) []string {
	h := strings.TrimSpace(r.Header.Get("Sec-Websocket-Protocol"))
	if h == "" {
		return nil
	}
	protocols := strings.Split(h, ",")
	for i := range protocols {
		protocols[i] = strings.TrimSpace(protocols[i])
	}
	return protocols
}
<|endoftext|>


<commit_before>package core

import (
	"flag"
	"net/http"
	"strconv"
)

var (
	// Production defines if the server must be using production settings.
	// It can be used by handlers to provide different logic for this environment.
	Production bool
	port       int
	beforeRun  []func()
)

func init() {
	flag.BoolVar(&Production, "production", false, "run the server in production environment")
	flag.IntVar(&port, "port", 8080, "the port to listen on")
	flag.Parse()
}

// BeforeRun adds a function that will be triggered before runnning the server.
func BeforeRun(f func()) {
	beforeRun = append(beforeRun, f)
}

// Run starts the server for listening and serving.
func Run() {
	if handlers == nil {
		panic("core: the handlers stack cannot be empty")
	}

	// Add a last handler to prevent "index out of range" errors if the previous last handler in stack calls Next().
	Use(func(c *Context) {
		http.Error(c.ResponseWriter, http.StatusText(http.StatusNotFound), http.StatusNotFound)
	})

	for _, f := range beforeRun {
		f()
	}

	panic(http.ListenAndServe(":"+strconv.Itoa(port), handlers))
}
<commit_msg>Use the standard NotFound<commit_after>package core

import (
	"flag"
	"net/http"
	"strconv"
)

var (
	// Production defines if the server must be using production settings.
	// It can be used by handlers to provide different logic for this environment.
	Production bool
	port       int
	beforeRun  []func()
)

func init() {
	flag.BoolVar(&Production, "production", false, "run the server in production environment")
	flag.IntVar(&port, "port", 8080, "the port to listen on")
	flag.Parse()
}

// BeforeRun adds a function that will be triggered before runnning the server.
func BeforeRun(f func()) {
	beforeRun = append(beforeRun, f)
}

// Run starts the server for listening and serving.
func Run() {
	if handlers == nil {
		panic("core: the handlers stack cannot be empty")
	}

	// Add a last handler to prevent "index out of range" errors if the previous last handler in stack calls Next().
	Use(func(c *Context) {
		http.NotFound(c.ResponseWriter, c.Request)
	})

	for _, f := range beforeRun {
		f()
	}

	panic(http.ListenAndServe(":"+strconv.Itoa(port), handlers))
}
<|endoftext|>


<commit_before>package xweb

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"net/http/pprof"
	"os"
	"runtime"
	runtimePprof "runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/go-xweb/httpsession"
	"github.com/go-xweb/log"
)

// ServerConfig is configuration for server objects.
type ServerConfig struct {
	Addr                   string
	Port                   int
	RecoverPanic           bool
	Profiler               bool
	EnableGzip             bool
	StaticExtensionsToGzip []string
	Url                    string
	UrlPrefix              string
	UrlSuffix              string
	StaticHtmlDir          string
	SessionTimeout         time.Duration
}

var ServerNumber uint = 0

// Server represents a xweb server.
type Server struct {
	Config         *ServerConfig
	Apps           map[string]*App
	AppName        map[string]string //[SWH|+]
	Name           string            //[SWH|+]
	SessionManager *httpsession.Manager
	RootApp        *App
	Logger         *log.Logger
	Env            map[string]interface{}
	//save the listener so it can be closed
	l              net.Listener
}

func NewServer(args ...string) *Server {
	name := ""
	if len(args) == 1 {
		name = args[0]
	} else {
		name = fmt.Sprintf("Server%d", ServerNumber)
		ServerNumber++
	}
	s := &Server{
		Config:  Config,
		Env:     map[string]interface{}{},
		Apps:    map[string]*App{},
		AppName: map[string]string{},
		Name:    name,
	}
	Servers[s.Name] = s //[SWH|+]

	s.SetLogger(log.New(os.Stdout, "", log.Ldefault()))

	app := NewApp("/", "root") //[SWH|+] ,"root"
	s.AddApp(app)
	return s
}

func (s *Server) AddApp(a *App) {
	a.BasePath = strings.TrimRight(a.BasePath, "/") + "/"
	s.Apps[a.BasePath] = a

	//[SWH|+]
	if a.Name != "" {
		s.AppName[a.Name] = a.BasePath
	}

	a.Server = s
	a.Logger = s.Logger
	if a.BasePath == "/" {
		s.RootApp = a
	}
}

func (s *Server) AddAction(cs ...interface{}) {
	s.RootApp.AddAction(cs...)
}

func (s *Server) AutoAction(c ...interface{}) {
	s.RootApp.AutoAction(c...)
}

func (s *Server) AddRouter(url string, c interface{}) {
	s.RootApp.AddRouter(url, c)
}

func (s *Server) AddTmplVar(name string, varOrFun interface{}) {
	s.RootApp.AddTmplVar(name, varOrFun)
}

func (s *Server) AddTmplVars(t *T) {
	s.RootApp.AddTmplVars(t)
}

func (s *Server) AddFilter(filter Filter) {
	s.RootApp.AddFilter(filter)
}

func (s *Server) AddConfig(name string, value interface{}) {
	s.RootApp.SetConfig(name, value)
}

func (s *Server) SetConfig(name string, value interface{}) {
	s.RootApp.SetConfig(name, value)
}

func (s *Server) GetConfig(name string) interface{} {
	return s.RootApp.GetConfig(name)
}

func (s *Server) error(w http.ResponseWriter, status int, content string) error {
	return s.RootApp.error(w, status, content)
}

func (s *Server) initServer() {
	if s.Config == nil {
		s.Config = &ServerConfig{}
		s.Config.Profiler = true
	}

	for _, app := range s.Apps {
		app.initApp()
	}
}

// ServeHTTP is the interface method for Go's http server package
func (s *Server) ServeHTTP(c http.ResponseWriter, req *http.Request) {
	s.Process(c, req)
}

// Process invokes the routing system for server s
// non-root app's route will override root app's if there is same path
func (s *Server) Process(w http.ResponseWriter, req *http.Request) {
	var result bool = true
	_, _ = XHook.Call("BeforeProcess", &result, s, w, req) //[SWH|+]call hook
	if !result {
		return
	}
	if s.Config.UrlSuffix != "" && strings.HasSuffix(req.URL.Path, s.Config.UrlSuffix) {
		req.URL.Path = strings.TrimSuffix(req.URL.Path, s.Config.UrlSuffix)
	}
	if s.Config.UrlPrefix != "" && strings.HasPrefix(req.URL.Path, "/"+s.Config.UrlPrefix) {
		req.URL.Path = strings.TrimPrefix(req.URL.Path, "/"+s.Config.UrlPrefix)
	}
	if req.URL.Path[0] != '/' {
		req.URL.Path = "/" + req.URL.Path
	}
	for _, app := range s.Apps {
		if app != s.RootApp && strings.HasPrefix(req.URL.Path, app.BasePath) {
			app.routeHandler(req, w)
			return
		}
	}
	s.RootApp.routeHandler(req, w)
	_, _ = XHook.Call("AfterProcess", &result, s, w, req) //[SWH|+]call hook
}

// Run starts the web application and serves HTTP requests for s
func (s *Server) Run(addr string) {
	addrs := strings.Split(addr, ":")
	s.Config.Addr = addrs[0]
	s.Config.Port, _ = strconv.Atoi(addrs[1])

	s.initServer()

	mux := http.NewServeMux()
	if s.Config.Profiler {
		mux.Handle("/debug/pprof", http.HandlerFunc(pprof.Index))
		mux.Handle("/debug/pprof/heap", pprof.Handler("heap"))
		mux.Handle("/debug/pprof/block", pprof.Handler("block"))
		mux.Handle("/debug/pprof/goroutine", pprof.Handler("goroutine"))
		mux.Handle("/debug/pprof/threadcreate", pprof.Handler("threadcreate"))

		mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
		mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
		mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))

		mux.Handle("/debug/pprof/startcpuprof", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			StartCPUProfile()
		}))
		mux.Handle("/debug/pprof/stopcpuprof", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			StopCPUProfile()
		}))
		mux.Handle("/debug/pprof/memprof", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			runtime.GC()
			runtimePprof.WriteHeapProfile(rw)
		}))
		mux.Handle("/debug/pprof/gc", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			PrintGCSummary(rw)
		}))

	}
	//[SWH|+]call hook
	if c, err := XHook.Call("MuxHandle", mux); err == nil {
		if ret := XHook.Value(c, 0); ret != nil {
			mux = ret.(*http.ServeMux)
		}
	}
	mux.Handle("/", s)

	s.Logger.Infof("http server is listening %s", addr)

	l, err := net.Listen("tcp", addr)
	if err != nil {
		s.Logger.Error("ListenAndServe:", err)
	}
	s.l = l
	err = http.Serve(s.l, mux)
	s.l.Close()
}

// RunFcgi starts the web application and serves FastCGI requests for s.
func (s *Server) RunFcgi(addr string) {
	s.initServer()
	s.Logger.Infof("fcgi server is listening %s", addr)
	s.listenAndServeFcgi(addr)
}

// RunScgi starts the web application and serves SCGI requests for s.
func (s *Server) RunScgi(addr string) {
	s.initServer()
	s.Logger.Infof("scgi server is listening %s", addr)
	s.listenAndServeScgi(addr)
}

// RunTLS starts the web application and serves HTTPS requests for s.
func (s *Server) RunTLS(addr string, config *tls.Config) error {
	s.initServer()
	mux := http.NewServeMux()
	mux.Handle("/", s)
	l, err := tls.Listen("tcp", addr, config)
	if err != nil {
		s.Logger.Errorf("Listen: %v", err)
		return err
	}

	s.l = l

	s.Logger.Infof("https server is listening %s", addr)

	return http.Serve(s.l, mux)
}

// Close stops server s.
func (s *Server) Close() {
	if s.l != nil {
		s.l.Close()
	}
}

// SetLogger sets the logger for server s
func (s *Server) SetLogger(logger *log.Logger) {
	s.Logger = logger
	s.Logger.SetPrefix("[" + s.Name + "] ")
	if s.RootApp != nil {
		s.RootApp.Logger = s.Logger
	}
}

func (s *Server) InitSession() {
	s.SessionManager = httpsession.Default()
	if s.Config.SessionTimeout > time.Second {
		s.SessionManager.SetMaxAge(s.Config.SessionTimeout)
	}
	s.SessionManager.Run()
	if s.RootApp != nil {
		s.RootApp.SessionManager = s.SessionManager
	}
}

func (s *Server) SetTemplateDir(path string) {
	s.RootApp.SetTemplateDir(path)
}

func (s *Server) SetStaticDir(path string) {
	s.RootApp.SetStaticDir(path)
}

func (s *Server) App(name string) *App {
	path, ok := s.AppName[name]
	if ok {
		return s.Apps[path]
	}
	return nil
}
<commit_msg>improved<commit_after>package xweb

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"net/http/pprof"
	"os"
	"runtime"
	runtimePprof "runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/go-xweb/httpsession"
	"github.com/go-xweb/log"
)

// ServerConfig is configuration for server objects.
type ServerConfig struct {
	Addr                   string
	Port                   int
	RecoverPanic           bool
	Profiler               bool
	EnableGzip             bool
	StaticExtensionsToGzip []string
	Url                    string
	UrlPrefix              string
	UrlSuffix              string
	StaticHtmlDir          string
	SessionTimeout         time.Duration
}

var ServerNumber uint = 0

// Server represents a xweb server.
type Server struct {
	Config         *ServerConfig
	Apps           map[string]*App
	AppName        map[string]string //[SWH|+]
	Name           string            //[SWH|+]
	SessionManager *httpsession.Manager
	RootApp        *App
	Logger         *log.Logger
	Env            map[string]interface{}
	//save the listener so it can be closed
	l net.Listener
}

func NewServer(args ...string) *Server {
	name := ""
	if len(args) == 1 {
		name = args[0]
	} else {
		name = fmt.Sprintf("Server%d", ServerNumber)
		ServerNumber++
	}
	s := &Server{
		Config:  Config,
		Env:     map[string]interface{}{},
		Apps:    map[string]*App{},
		AppName: map[string]string{},
		Name:    name,
	}
	Servers[s.Name] = s //[SWH|+]

	s.SetLogger(log.New(os.Stdout, "", log.Ldefault()))

	app := NewApp("/", "root") //[SWH|+] ,"root"
	s.AddApp(app)
	return s
}

func (s *Server) AddApp(a *App) {
	a.BasePath = strings.TrimRight(a.BasePath, "/") + "/"
	s.Apps[a.BasePath] = a

	//[SWH|+]
	if a.Name != "" {
		s.AppName[a.Name] = a.BasePath
	}

	a.Server = s
	a.Logger = s.Logger
	if a.BasePath == "/" {
		s.RootApp = a
	}
}

func (s *Server) AddAction(cs ...interface{}) {
	s.RootApp.AddAction(cs...)
}

func (s *Server) AutoAction(c ...interface{}) {
	s.RootApp.AutoAction(c...)
}

func (s *Server) AddRouter(url string, c interface{}) {
	s.RootApp.AddRouter(url, c)
}

func (s *Server) AddTmplVar(name string, varOrFun interface{}) {
	s.RootApp.AddTmplVar(name, varOrFun)
}

func (s *Server) AddTmplVars(t *T) {
	s.RootApp.AddTmplVars(t)
}

func (s *Server) AddFilter(filter Filter) {
	s.RootApp.AddFilter(filter)
}

func (s *Server) AddConfig(name string, value interface{}) {
	s.RootApp.SetConfig(name, value)
}

func (s *Server) SetConfig(name string, value interface{}) {
	s.RootApp.SetConfig(name, value)
}

func (s *Server) GetConfig(name string) interface{} {
	return s.RootApp.GetConfig(name)
}

func (s *Server) error(w http.ResponseWriter, status int, content string) error {
	return s.RootApp.error(w, status, content)
}

func (s *Server) initServer() {
	if s.Config == nil {
		s.Config = &ServerConfig{}
		s.Config.Profiler = true
	}

	for _, app := range s.Apps {
		app.initApp()
	}
}

// ServeHTTP is the interface method for Go's http server package
func (s *Server) ServeHTTP(c http.ResponseWriter, req *http.Request) {
	s.Process(c, req)
}

// Process invokes the routing system for server s
// non-root app's route will override root app's if there is same path
func (s *Server) Process(w http.ResponseWriter, req *http.Request) {
	var result bool = true
	_, _ = XHook.Call("BeforeProcess", &result, s, w, req) //[SWH|+]call hook
	if !result {
		return
	}
	if s.Config.UrlSuffix != "" && strings.HasSuffix(req.URL.Path, s.Config.UrlSuffix) {
		req.URL.Path = strings.TrimSuffix(req.URL.Path, s.Config.UrlSuffix)
	}
	if s.Config.UrlPrefix != "" && strings.HasPrefix(req.URL.Path, "/"+s.Config.UrlPrefix) {
		req.URL.Path = strings.TrimPrefix(req.URL.Path, "/"+s.Config.UrlPrefix)
	}
	if req.URL.Path[0] != '/' {
		req.URL.Path = "/" + req.URL.Path
	}
	for _, app := range s.Apps {
		if app != s.RootApp && strings.HasPrefix(req.URL.Path, app.BasePath) {
			app.routeHandler(req, w)
			return
		}
	}
	s.RootApp.routeHandler(req, w)
	_, _ = XHook.Call("AfterProcess", &result, s, w, req) //[SWH|+]call hook
}

// Run starts the web application and serves HTTP requests for s
func (s *Server) Run(addr string) {
	addrs := strings.Split(addr, ":")
	s.Config.Addr = addrs[0]
	s.Config.Port, _ = strconv.Atoi(addrs[1])

	s.initServer()

	mux := http.NewServeMux()
	if s.Config.Profiler {
		mux.Handle("/debug/pprof", http.HandlerFunc(pprof.Index))
		mux.Handle("/debug/pprof/heap", pprof.Handler("heap"))
		mux.Handle("/debug/pprof/block", pprof.Handler("block"))
		mux.Handle("/debug/pprof/goroutine", pprof.Handler("goroutine"))
		mux.Handle("/debug/pprof/threadcreate", pprof.Handler("threadcreate"))

		mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
		mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
		mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))

		mux.Handle("/debug/pprof/startcpuprof", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			StartCPUProfile()
		}))
		mux.Handle("/debug/pprof/stopcpuprof", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			StopCPUProfile()
		}))
		mux.Handle("/debug/pprof/memprof", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			runtime.GC()
			runtimePprof.WriteHeapProfile(rw)
		}))
		mux.Handle("/debug/pprof/gc", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			PrintGCSummary(rw)
		}))

	}
	//[SWH|+]call hook
	if c, err := XHook.Call("MuxHandle", mux); err == nil {
		if ret := XHook.Value(c, 0); ret != nil {
			mux = ret.(*http.ServeMux)
		}
	}
	mux.Handle("/", s)

	s.Logger.Infof("http server is listening %s", addr)

	l, err := net.Listen("tcp", addr)
	if err != nil {
		s.Logger.Error("ListenAndServe:", err)
	}
	s.l = l
	err = http.Serve(s.l, mux)
	s.l.Close()
}

// RunFcgi starts the web application and serves FastCGI requests for s.
func (s *Server) RunFcgi(addr string) {
	s.initServer()
	s.Logger.Infof("fcgi server is listening %s", addr)
	s.listenAndServeFcgi(addr)
}

// RunScgi starts the web application and serves SCGI requests for s.
func (s *Server) RunScgi(addr string) {
	s.initServer()
	s.Logger.Infof("scgi server is listening %s", addr)
	s.listenAndServeScgi(addr)
}

// RunTLS starts the web application and serves HTTPS requests for s.
func (s *Server) RunTLS(addr string, config *tls.Config) error {
	s.initServer()
	mux := http.NewServeMux()
	mux.Handle("/", s)
	l, err := tls.Listen("tcp", addr, config)
	if err != nil {
		s.Logger.Errorf("Listen: %v", err)
		return err
	}

	s.l = l

	s.Logger.Infof("https server is listening %s", addr)

	return http.Serve(s.l, mux)
}

// Close stops server s.
func (s *Server) Close() {
	if s.l != nil {
		s.l.Close()
	}
}

// SetLogger sets the logger for server s
func (s *Server) SetLogger(logger *log.Logger) {
	s.Logger = logger
	s.Logger.SetPrefix("[" + s.Name + "] ")
	if s.RootApp != nil {
		s.RootApp.Logger = s.Logger
	}
}

func (s *Server) InitSession() {
	if s.SessionManager == nil {
		s.SessionManager = httpsession.Default()
	}
	if s.Config.SessionTimeout > time.Second {
		s.SessionManager.SetMaxAge(s.Config.SessionTimeout)
	}
	s.SessionManager.Run()
	if s.RootApp != nil {
		s.RootApp.SessionManager = s.SessionManager
	}
}

func (s *Server) SetTemplateDir(path string) {
	s.RootApp.SetTemplateDir(path)
}

func (s *Server) SetStaticDir(path string) {
	s.RootApp.SetStaticDir(path)
}

func (s *Server) App(name string) *App {
	path, ok := s.AppName[name]
	if ok {
		return s.Apps[path]
	}
	return nil
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

package servicelib

// Search types and various encapsulation used in certain API responses

import (
	"github.com/jvehent/gozdef"
)

// Parameters used for a search request.
type SearchParams struct {
	Searches []Search `json:"search"` // Slice of searches to conduct for this request
}

// Describes an individual search.
type Search struct {
	Identifier string `json:"identifier,omitempty"` // A unique element identifier
	Host       string `json:"host,omitempty"`       // Host search criteria
	Confidence int    `json:"confidence,omitempty"` // Reporter confidence
}

// Describes a response to a search request.
type SearchResponse struct {
	SearchID string `json:"id"` // The ID of the search to fetch results
}

// The response to a search results request.
type SearchIDResponse struct {
	Results []SearchResult `json:"results"` // A slice of results for the search
}

// An individual search result.
type SearchResult struct {
	Identifier string  `json:"identifier"` // Search request element identifier (from Search)
	Service    Service `json:"service"`    // The service result information
}

// Parameters to an indicator request
type IndicatorParams struct {
	Indicators []Indicator `json:"indicators"` // Slice of indicators to process
}

// Describes indicator information
type Indicator struct {
	Host      string `json:"host,omitempty"`
	Class     string `json:"class"`
	CheckType string `json:"checktype"`
	Status    bool   `json:"status"`
}

// The response to an indicator request
type IndicatorResponse struct {
	OK bool `json:"ok"`
}

type RRAUpdateResponse struct {
	OK bool `json:"ok"`
}

// The response to a general service query.
type Service struct {
	Services    []RRAService `json:"services,omitempty"`    // Services linked from RRA table
	SystemGroup SystemGroup  `json:"systemgroup,omitempty"` // Database system group
	Owner       Owner        `json:"owner,omitempty"`       // Ownership details
	Found       bool         `json:"found"`                 // Results of search
}

// The response to a system group list request.
type SystemGroupsResponse struct {
	Results []SystemGroup `json:"results"`
}

// The response to an RRA list request.
type RRAsResponse struct {
	Results []RRAService `json:"results"`
}

// The response to a match search request.
type SearchMatchResponse struct {
	Hosts []Host `json:"hosts"`
}

// Response to vulnerabilities request for a target
type VulnsTargetResponse struct {
	Vulnerabilities []gozdef.VulnEvent `json:"vulnerabilities"`
}

// Response to /risks request, includes all known RRAs
type RisksResponse struct {
	Risks []RRAServiceRisk `json:"risks"`
}
<commit_msg>update indicators to include mig statistics<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

package servicelib

// Search types and various encapsulation used in certain API responses

import (
	"github.com/jvehent/gozdef"
)

// Parameters used for a search request.
type SearchParams struct {
	Searches []Search `json:"search"` // Slice of searches to conduct for this request
}

// Describes an individual search.
type Search struct {
	Identifier string `json:"identifier,omitempty"` // A unique element identifier
	Host       string `json:"host,omitempty"`       // Host search criteria
	Confidence int    `json:"confidence,omitempty"` // Reporter confidence
}

// Describes a response to a search request.
type SearchResponse struct {
	SearchID string `json:"id"` // The ID of the search to fetch results
}

// The response to a search results request.
type SearchIDResponse struct {
	Results []SearchResult `json:"results"` // A slice of results for the search
}

// An individual search result.
type SearchResult struct {
	Identifier string  `json:"identifier"` // Search request element identifier (from Search)
	Service    Service `json:"service"`    // The service result information
}

// Parameters to an indicator request
type IndicatorParams struct {
	Indicators []Indicator `json:"indicators"` // Slice of indicators to process
}

// Describes indicator information
type Indicator struct {
	// Common attributes amongst all classes
	Host string `json:"host,omitempty"` // Hostname of system

	Class string `json:"class"` // Indicator class

	// Used for "vuln" class indicators
	CheckType string `json:"checktype"`
	Status    bool   `json:"status"`

	// Used for "migagent" class indicators
	MIG IndicatorMIG `json:"mig"`
}

type IndicatorMIG struct {
	MIGHostname string      `json:"hostname"`
	MIGVersion  string      `json:"version"`
	Tags        interface{} `json:"tags"`
	Environment interface{} `json:"environment"`
}

// The response to an indicator request
type IndicatorResponse struct {
	OK bool `json:"ok"`
}

type RRAUpdateResponse struct {
	OK bool `json:"ok"`
}

// The response to a general service query.
type Service struct {
	Services    []RRAService `json:"services,omitempty"`    // Services linked from RRA table
	SystemGroup SystemGroup  `json:"systemgroup,omitempty"` // Database system group
	Owner       Owner        `json:"owner,omitempty"`       // Ownership details
	Found       bool         `json:"found"`                 // Results of search
}

// The response to a system group list request.
type SystemGroupsResponse struct {
	Results []SystemGroup `json:"results"`
}

// The response to an RRA list request.
type RRAsResponse struct {
	Results []RRAService `json:"results"`
}

// The response to a match search request.
type SearchMatchResponse struct {
	Hosts []Host `json:"hosts"`
}

// Response to vulnerabilities request for a target
type VulnsTargetResponse struct {
	Vulnerabilities []gozdef.VulnEvent `json:"vulnerabilities"`
}

// Response to /risks request, includes all known RRAs
type RisksResponse struct {
	Risks []RRAServiceRisk `json:"risks"`
}
<|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"
	"sort"
	"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)
	}
}

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, " ")
	groups := []string{}
	if len(form.Groups) != 0 {
		groups = append(groups, strings.Split(form.Groups, ",")...)
	}
	sort.Strings(groups)

	encoded := encodeString(form.Domain)
	encodedURL := encodeString(form.URL)

	c, err := redisPool.Get()
	errHndlr(err)
	defer redisPool.Put(c)
	pipedCommands := 0
	for _, prefix := range getPrefixes(form.Title) {
		if len(groups) == 0 {
			c.Append("ZADD", encoded+prefix, form.Popularity, encodedURL)
			pipedCommands++
		}
		for _, group := range groups {
			encodedGroup := encodeString(group)
			c.Append("ZADD", encoded+encodedGroup+prefix, form.Popularity, encodedURL)
			pipedCommands++
		}
	}
	c.Append("HSET", encoded+"$titles", encodedURL, form.Title)
	pipedCommands++
	c.Append("HSET", encoded+"$urls", encodedURL, form.URL)
	pipedCommands++
	for i := 1; i <= pipedCommands; i++ {
		if err := c.GetReply().Err; err != nil {
			errHndlr(err)
		}
	}

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

type deleteForm struct {
	Domain string
	URL    string
}

// FieldMap defines the bindings for deleteForm
func (f *deleteForm) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&f.Domain: binding.Field{
			Form:     "domain",
			Required: true,
		},
		&f.URL: binding.Field{
			Form:     "url",
			Required: true,
		},
	}
}

// Validate checks the values for deleteForm
func (f deleteForm) 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.URL, " ") == "" {
		errs = append(errs, binding.Error{
			FieldNames:     []string{"url"},
			Classification: "ComplaintError",
			Message:        "Can't be empty",
		})
	}
	return errs
}

func deleteHandler(w http.ResponseWriter, req *http.Request) {
	form := new(deleteForm)
	errs := binding.Bind(req, form)
	if errs.Handle(w) {
		return
	}
	form.Domain = strings.Trim(form.Domain, " ")
	form.URL = strings.Trim(form.URL, " ")

	encoded := encodeString(form.Domain)

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

	encodedURL := encodeString(form.URL)
	var title string
	title, err = c.Cmd("HGET", encoded+"$titles", encodedURL).Str()
	if err != nil {
		errHndlr(err)
	}
	prefixes := getPrefixes(title)
	pipedCommands := 0
	for _, prefix := range prefixes {
		c.Append("ZREM", encoded+prefix, encodedURL)
		pipedCommands++
	}

	c.Append("HDEL", encoded+"$titles", encodedURL)
	pipedCommands++
	c.Append("HDEL", encoded+"$urls", encodedURL)
	pipedCommands++

	for i := 1; i <= pipedCommands; i++ {
		if err := c.GetReply().Err; err != nil {
			errHndlr(err)
		}
	}
	output := map[string]string{"message": "OK"}
	renderer.JSON(w, http.StatusNoContent, output)
}

type fetchForm struct {
	Number int
	Groups string
	Query  string
	Domain string
}

// FieldMap defines the bindings for fetchForm
func (f *fetchForm) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&f.Number: "n",
		&f.Groups: "g",
		&f.Query: binding.Field{
			Form:     "q",
			Required: true,
		},
		&f.Domain: binding.Field{
			Form:     "d",
			Required: true,
		},
	}
}

// Reply isn't a great name :)
type Reply struct {
	URL   string
	Score string
}

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, " ")
	var groups = make([]string, strings.Count(form.Groups, ",")+1)
	for i, piece := range strings.Split(form.Groups, ",") {
		groups[i] = strings.Trim(piece, " ")
	}
	sort.Strings(groups)

	encoded := encodeString(form.Domain)

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

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

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

		c.GetReply() // the ZINTERSTORE
		replies, err := c.GetReply().List()
		return replies, err
	}

	replies, err := getReplies(terms, "")
	errHndlr(err)
	var replyStructs []Reply
	for i, element := range replies {
		if i%2 == 0 {
			replyStructs = append(replyStructs, Reply{element, replies[i+1]})
		}
	}
	for _, group := range groups {
		replies, err = getReplies(terms, group)
		errHndlr(err)
		for i, element := range replies {
			if i%2 == 0 {
				replyStructs = append(replyStructs, Reply{element, replies[i+1]})
			}
		}
	}
	RemoveDuplicates := func(xs *[]Reply) {
		found := make(map[string]bool)
		j := 0
		for i, x := range *xs {
			if !found[x.URL] {
				found[x.URL] = true
				(*xs)[j] = (*xs)[i]
				j++
			}
		}
		*xs = (*xs)[:j]
	}
	RemoveDuplicates(&replyStructs)
	// We might want to sort this here by the extra business logic
	// on sorting.
	encodedUrls := make([]string, len(replyStructs))
	for i, each := range replyStructs {
		encodedUrls[i] = each.URL
	}

	var titles []string
	var urls []string
	if len(encodedUrls) == 0 {
	} else {
		titles, err = c.Cmd("HMGET", encoded+"$titles", encodedUrls).List()
		errHndlr(err)
		urls, err = c.Cmd("HMGET", encoded+"$urls", encodedUrls).List()
		errHndlr(err)
	}
	rows := make([]interface{}, len(titles))
	for i, title := range titles {
		row := make([]string, 2)
		row[0] = urls[i]
		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 (
	redisPool *pool.Pool
	procs     int
	debug     = true
	renderer  = render.New()
	redisURL  = "127.0.0.1:6379"
)

func main() {
	var (
		port          = 3001
		redisDatabase = 0
		redisPoolSize = 10
	)
	flag.IntVar(&port, "port", port, "Port to start the server on")
	flag.IntVar(&procs, "procs", 1, "Number of CPU processors (0 to use max)")
	flag.BoolVar(&debug, "debug", false, "Debug mode")
	flag.StringVar(
		&redisURL, "redisURL", redisURL,
		"Redis URL to tcp connect to")
	flag.IntVar(&redisDatabase, "redisDatabase", redisDatabase,
		"Redis database number to connect to")
	flag.Parse()

	fmt.Println("REDIS DATABASE:", redisDatabase)
	fmt.Println("DEBUG MODE:", debug)

	if !debug {
		redisPoolSize = 100
	}

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

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

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

	var err error
	redisPool, err = pool.NewCustomPool("tcp", redisURL, redisPoolSize, df)
	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")
	mux.HandleFunc("/v1", deleteHandler).Methods("DELETE")

	n := negroni.Classic()
	n.UseHandler(mux)
	n.Run(fmt.Sprintf(":%d", port))
}
<commit_msg>only use ZINTERSTORE on multiple terms<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"
	"regexp"
	"runtime"
	"sort"
	"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)
	}
}

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, " ")
	groups := []string{}
	if len(form.Groups) != 0 {
		groups = append(groups, strings.Split(form.Groups, ",")...)
	}
	sort.Strings(groups)

	encoded := encodeString(form.Domain)
	encodedURL := encodeString(form.URL)

	c, err := redisPool.Get()
	errHndlr(err)
	defer redisPool.Put(c)
	pipedCommands := 0
	for _, prefix := range getPrefixes(form.Title) {
		if len(groups) == 0 {
			c.Append("ZADD", encoded+prefix, form.Popularity, encodedURL)
			pipedCommands++
		}
		for _, group := range groups {
			encodedGroup := encodeString(group)
			c.Append("ZADD", encoded+encodedGroup+prefix, form.Popularity, encodedURL)
			pipedCommands++
		}
	}
	c.Append("HSET", encoded+"$titles", encodedURL, form.Title)
	pipedCommands++
	c.Append("HSET", encoded+"$urls", encodedURL, form.URL)
	pipedCommands++
	for i := 1; i <= pipedCommands; i++ {
		if err := c.GetReply().Err; err != nil {
			errHndlr(err)
		}
	}

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

type deleteForm struct {
	Domain string
	URL    string
}

// FieldMap defines the bindings for deleteForm
func (f *deleteForm) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&f.Domain: binding.Field{
			Form:     "domain",
			Required: true,
		},
		&f.URL: binding.Field{
			Form:     "url",
			Required: true,
		},
	}
}

// Validate checks the values for deleteForm
func (f deleteForm) 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.URL, " ") == "" {
		errs = append(errs, binding.Error{
			FieldNames:     []string{"url"},
			Classification: "ComplaintError",
			Message:        "Can't be empty",
		})
	}
	return errs
}

func deleteHandler(w http.ResponseWriter, req *http.Request) {
	form := new(deleteForm)
	errs := binding.Bind(req, form)
	if errs.Handle(w) {
		return
	}
	form.Domain = strings.Trim(form.Domain, " ")
	form.URL = strings.Trim(form.URL, " ")

	encoded := encodeString(form.Domain)

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

	encodedURL := encodeString(form.URL)
	var title string
	title, err = c.Cmd("HGET", encoded+"$titles", encodedURL).Str()
	if err != nil {
		errHndlr(err)
	}
	prefixes := getPrefixes(title)
	pipedCommands := 0
	for _, prefix := range prefixes {
		c.Append("ZREM", encoded+prefix, encodedURL)
		pipedCommands++
	}

	c.Append("HDEL", encoded+"$titles", encodedURL)
	pipedCommands++
	c.Append("HDEL", encoded+"$urls", encodedURL)
	pipedCommands++

	for i := 1; i <= pipedCommands; i++ {
		if err := c.GetReply().Err; err != nil {
			errHndlr(err)
		}
	}
	output := map[string]string{"message": "OK"}
	renderer.JSON(w, http.StatusNoContent, output)
}

type fetchForm struct {
	Number int
	Groups string
	Query  string
	Domain string
}

// FieldMap defines the bindings for fetchForm
func (f *fetchForm) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&f.Number: "n",
		&f.Groups: "g",
		&f.Query: binding.Field{
			Form:     "q",
			Required: true,
		},
		&f.Domain: binding.Field{
			Form:     "d",
			Required: true,
		},
	}
}

// Reply isn't a great name :)
type Reply struct {
	URL   string
	Score string
}

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, " ")

	groups := []string{}
	if len(form.Groups) != 0 {
		groups = append(groups, strings.Split(form.Groups, ",")...)
	}
	sort.Strings(groups)

	encoded := encodeString(form.Domain)

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

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

	getReplies := func(terms []string, group string) ([]string, error) {
		encodedTerms := make([]string, len(terms))
		encodedGroup := ""
		if group != "" {
			encodedGroup = encodeString(group)
		}
		for i, term := range terms {
			encodedTerms[i] = encoded + encodedGroup + term
		}
		var replies []string
		if len(terms) > 1 {
			c.Append("ZINTERSTORE", "$tmp", len(terms), encodedTerms, "AGGREGATE", "max")
			c.Append("ZREVRANGE", "$tmp", 0, n-1, "WITHSCORES")
			c.GetReply() // the ZINTERSTORE
			replies, err = c.GetReply().List()
		} else {
			replies, err = c.Cmd("ZREVRANGE", encodedTerms[0], 0, n-1, "WITHSCORES").List()
		}
		return replies, err
	}

	replies, err := getReplies(terms, "")
	errHndlr(err)
	var replyStructs []Reply
	for i, element := range replies {
		if i%2 == 0 {
			replyStructs = append(replyStructs, Reply{element, replies[i+1]})
		}
	}
	for _, group := range groups {
		replies, err = getReplies(terms, group)
		errHndlr(err)
		for i, element := range replies {
			if i%2 == 0 {
				replyStructs = append(replyStructs, Reply{element, replies[i+1]})
			}
		}
	}
	RemoveDuplicates := func(xs *[]Reply) {
		found := make(map[string]bool)
		j := 0
		for i, x := range *xs {
			if !found[x.URL] {
				found[x.URL] = true
				(*xs)[j] = (*xs)[i]
				j++
			}
		}
		*xs = (*xs)[:j]
	}
	RemoveDuplicates(&replyStructs)
	// We might want to sort this here by the extra business logic
	// on sorting.
	encodedUrls := make([]string, len(replyStructs))
	for i, each := range replyStructs {
		encodedUrls[i] = each.URL
	}

	var titles []string
	var urls []string
	if len(encodedUrls) == 0 {
	} else {
		titles, err = c.Cmd("HMGET", encoded+"$titles", encodedUrls).List()
		errHndlr(err)
		urls, err = c.Cmd("HMGET", encoded+"$urls", encodedUrls).List()
		errHndlr(err)
	}
	rows := make([]interface{}, len(titles))
	for i, title := range titles {
		row := make([]string, 2)
		row[0] = urls[i]
		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 (
	redisPool *pool.Pool
	procs     int
	debug     = true
	renderer  = render.New()
	redisURL  = "127.0.0.1:6379"
)

func main() {
	var (
		port          = 3001
		redisDatabase = 0
		redisPoolSize = 10
	)
	flag.IntVar(&port, "port", port, "Port to start the server on")
	flag.IntVar(&procs, "procs", 1, "Number of CPU processors (0 to use max)")
	flag.BoolVar(&debug, "debug", false, "Debug mode")
	flag.StringVar(
		&redisURL, "redisURL", redisURL,
		"Redis URL to tcp connect to")
	flag.IntVar(&redisDatabase, "redisDatabase", redisDatabase,
		"Redis database number to connect to")
	flag.Parse()

	fmt.Println("REDIS DATABASE:", redisDatabase)
	fmt.Println("DEBUG MODE:", debug)

	if !debug {
		redisPoolSize = 100
	}

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

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

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

	var err error
	redisPool, err = pool.NewCustomPool("tcp", redisURL, redisPoolSize, df)
	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")
	mux.HandleFunc("/v1", deleteHandler).Methods("DELETE")

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


<commit_before>package main

import (
	"code.google.com/p/go.net/websocket"
	"github.com/kr/pty"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var ide *template.Template
var api *template.Template
var hostname, _ = ioutil.ReadFile("/etc/hostname")
var users = make(map[string]string)
var config = map[string]string{"port": "80", "projectdir": "/projects/", "WebviewUrl": "/website/"}

func main() {
	content, err := ioutil.ReadFile("/etc/ide/ide.html") //ide.html is actually a go template
	acontent, aerr := ioutil.ReadFile("/etc/ide/api.html")
	settings, serr := ioutil.ReadFile("/etc/ide/settings.conf")

	if err != nil {
		panic(err)
	}
	if aerr != nil {
		panic(aerr)
	}
	if serr == nil {
		set := strings.Split(string(settings), "\n")
		for _, line := range set {
			if len(line) > 0 && line[0:1] != "#" {
				line = strings.Split(line, "#")[0]
				lsplit := strings.Split(line, " ")
				config[strings.ToLower(lsplit[0])] = strings.TrimRight(line[len(lsplit[0])+1:], " ")
			}
		}
	}
	os.Mkdir(config["projectdir"], 0775)
	ide, err = template.New("page").Parse(string(content))
	api, err = template.New("page").Parse(string(acontent))
	if err != nil {
		panic(err)
	}
	http.HandleFunc("/", index) //All requests to / and 404s will route to Index
	http.HandleFunc("/ide.js", ideJs)
	http.HandleFunc("/cm.js", cmJs)
	http.HandleFunc("/cm.css", cmCss)
	http.HandleFunc("/python.js", pythonMode)
	http.HandleFunc("/shell.js", shellMode)
	http.HandleFunc("/api/", apiDoc)
	http.HandleFunc("/api/listfiles", listFiles)
	http.HandleFunc("/api/listthemes", listThemes)
	http.HandleFunc("/api/usernumber", userNumber)
	http.HandleFunc("/api/userleave", userLeave)
	http.HandleFunc("/api/readfile", readFile)
	http.HandleFunc("/api/savefile", saveFile)
	http.HandleFunc("/api/hostname", hostnameOut)
	http.Handle("/api/socket", websocket.Handler(socketServer))
	http.Handle("/website/", http.StripPrefix("/website/", http.FileServer(http.Dir("/etc/ide/website"))))
	http.Handle("/images/", http.StripPrefix("/images/", http.FileServer(http.Dir("/etc/ide/assets/images"))))
	http.Handle("/themes/", http.StripPrefix("/themes/", http.FileServer(http.Dir("/etc/ide/assets/themes"))))
	err = http.ListenAndServe(":"+config["port"], nil)
	if err != nil {
		panic(err)
	}
}

//Called by requests to / and 404s
func index(w http.ResponseWriter, r *http.Request) {
	data := map[string]string{
		"title":   "RStem IDE@" + string(hostname),
		"baseurl": config["WebviewUrl"],
	}
	ide.ExecuteTemplate(w, "page", data)
}

func apiDoc(w http.ResponseWriter, r *http.Request) {
	api.ExecuteTemplate(w, "page", nil)
}

func ideJs(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/ide.js")
}

func cmJs(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/cmirror/codemirror.js")
}

func cmCss(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/cmirror/codemirror.css")
}

func pythonMode(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/cmirror/python.js")
}

func shellMode(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/cmirror/shell.js")
}

//Returns the number of users using a particular file
func userNumber(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	opts := r.URL.Query()
	fname := strings.Trim(opts.Get("file"), " ./")
	count := 0
	for _, v := range users {
		if v == fname {
			count++
		}
	}
	io.WriteString(w, strconv.Itoa(count))
}

func userLeave(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	users[strings.Split(r.RemoteAddr, ":")[0]] = "exited"
}

//Lists all files in the projects directory
func listFiles(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	files, _ := ioutil.ReadDir(config["projectdir"])
	var list string
	for _, f := range files {
		if !f.IsDir() {
			list += "\n" + f.Name()
		}
	}
	if list != "" {
		io.WriteString(w, list[1:])
	}
}

//Lists the availible themes
func listThemes(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	files, _ := ioutil.ReadDir("/etc/ide/assets/themes/")
	var list string
	for _, f := range files {
		if !f.IsDir() {
			list += "\n" + f.Name()
		}
	}
	if list != "" {
		io.WriteString(w, list[1:])
	}
}

//Outputs the contents of the requested file
func readFile(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	opts := r.URL.Query()
	fname := strings.Trim(opts.Get("file"), " ./")
	contents, err := ioutil.ReadFile(config["projectdir"] + fname)
	if err != nil {
		io.WriteString(w, "error")
		return
	}
	users[strings.Split(r.RemoteAddr, ":")[0]] = fname
	io.WriteString(w, string(contents))
}

//Potential issue here because POST requests tend to have size limits
func saveFile(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	name := strings.Trim(r.Form.Get("file"), " ./")
	content := r.Form.Get("content")
	ioutil.WriteFile(config["projectdir"]+name, []byte(content), 0744)
}

func hostnameOut(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	io.WriteString(w, string(hostname))
}

//Called as a goroutine to wait for the close command and kill the process.
func watchClose(s *websocket.Conn, p *os.Process) {
	data := make([]byte, 512)
	n, _ := s.Read(data)
	payload := string(data[:n]) //[:n] to cut out padding
	if payload == "close" {
		p.Signal(syscall.SIGINT)
	}
}

//Receives socket connections to /api/socket and creates a PTY to run the process
func socketServer(s *websocket.Conn) {
	data := make([]byte, 512)
	n, _ := s.Read(data)
	com := exec.Command(config["projectdir"] + string(data[:n])) //[:n] to cut out padding
	pt, err := pty.Start(com)
	if err != nil {
		s.Write([]byte("error: " + err.Error()))
		s.Close()
		return
	}
	go watchClose(s, com.Process)
	for {
		out := make([]byte, 1024)
		n, err := pt.Read(out)
		if err != nil {
			if err.Error() == "read /dev/ptmx: input/output error" || err.Error() == "EOF" {
				break
			}
			s.Write([]byte("error: " + err.Error()))
		} else if n > 0 {
			s.Write([]byte("output: " + string(out[:n])))
		}
		time.Sleep(time.Millisecond * 100) //For minimal CPU impact
	}
	s.Close()
}
<commit_msg>Started change relay socket<commit_after>package main

import (
	"code.google.com/p/go.net/websocket"
	"github.com/kr/pty"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var ide *template.Template
var api *template.Template
var hostname, _ = ioutil.ReadFile("/etc/hostname")
var users = make(map[string]string)
var config = map[string]string{"port": "80", "projectdir": "/projects/", "WebviewUrl": "/website/"}

func main() {
	content, err := ioutil.ReadFile("/etc/ide/ide.html") //ide.html is actually a go template
	acontent, aerr := ioutil.ReadFile("/etc/ide/api.html")
	settings, serr := ioutil.ReadFile("/etc/ide/settings.conf")

	if err != nil {
		panic(err)
	}
	if aerr != nil {
		panic(aerr)
	}
	if serr == nil {
		set := strings.Split(string(settings), "\n")
		for _, line := range set {
			if len(line) > 0 && line[0:1] != "#" {
				line = strings.Split(line, "#")[0]
				lsplit := strings.Split(line, " ")
				config[strings.ToLower(lsplit[0])] = strings.TrimRight(line[len(lsplit[0])+1:], " ")
			}
		}
	}
	os.Mkdir(config["projectdir"], 0775)
	ide, err = template.New("page").Parse(string(content))
	api, err = template.New("page").Parse(string(acontent))
	if err != nil {
		panic(err)
	}
	http.HandleFunc("/", index) //All requests to / and 404s will route to Index
	http.HandleFunc("/ide.js", ideJs)
	http.HandleFunc("/cm.js", cmJs)
	http.HandleFunc("/cm.css", cmCss)
	http.HandleFunc("/python.js", pythonMode)
	http.HandleFunc("/shell.js", shellMode)
	http.HandleFunc("/api/", apiDoc)
	http.HandleFunc("/api/listfiles", listFiles)
	http.HandleFunc("/api/listthemes", listThemes)
	http.HandleFunc("/api/usernumber", userNumber)
	http.HandleFunc("/api/userleave", userLeave)
	http.HandleFunc("/api/readfile", readFile)
	http.HandleFunc("/api/savefile", saveFile)
	http.HandleFunc("/api/hostname", hostnameOut)
	http.Handle("/api/socket", websocket.Handler(socketServer))
	http.Handle("/api/changes", websocket.Handler(changeServer))
	http.Handle("/website/", http.StripPrefix("/website/", http.FileServer(http.Dir("/etc/ide/website"))))
	http.Handle("/images/", http.StripPrefix("/images/", http.FileServer(http.Dir("/etc/ide/assets/images"))))
	http.Handle("/themes/", http.StripPrefix("/themes/", http.FileServer(http.Dir("/etc/ide/assets/themes"))))
	err = http.ListenAndServe(":"+config["port"], nil)
	if err != nil {
		panic(err)
	}
}

//Called by requests to / and 404s
func index(w http.ResponseWriter, r *http.Request) {
	data := map[string]string{
		"title":   "RStem IDE@" + string(hostname),
		"baseurl": config["WebviewUrl"],
	}
	ide.ExecuteTemplate(w, "page", data)
}

func apiDoc(w http.ResponseWriter, r *http.Request) {
	api.ExecuteTemplate(w, "page", nil)
}

func ideJs(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/ide.js")
}

func cmJs(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/cmirror/codemirror.js")
}

func cmCss(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/cmirror/codemirror.css")
}

func pythonMode(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/cmirror/python.js")
}

func shellMode(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "/etc/ide/assets/cmirror/shell.js")
}

//Returns the number of users using a particular file
func userNumber(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	opts := r.URL.Query()
	fname := strings.Trim(opts.Get("file"), " ./")
	count := 0
	for _, v := range users {
		if v == fname {
			count++
		}
	}
	io.WriteString(w, strconv.Itoa(count))
}

func userLeave(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	users[strings.Split(r.RemoteAddr, ":")[0]] = "exited"
}

//Lists all files in the projects directory
func listFiles(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	files, _ := ioutil.ReadDir(config["projectdir"])
	var list string
	for _, f := range files {
		if !f.IsDir() {
			list += "\n" + f.Name()
		}
	}
	if list != "" {
		io.WriteString(w, list[1:])
	}
}

//Lists the availible themes
func listThemes(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	files, _ := ioutil.ReadDir("/etc/ide/assets/themes/")
	var list string
	for _, f := range files {
		if !f.IsDir() {
			list += "\n" + f.Name()
		}
	}
	if list != "" {
		io.WriteString(w, list[1:])
	}
}

//Outputs the contents of the requested file
func readFile(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	opts := r.URL.Query()
	fname := strings.Trim(opts.Get("file"), " ./")
	contents, err := ioutil.ReadFile(config["projectdir"] + fname)
	if err != nil {
		io.WriteString(w, "error")
		return
	}
	users[strings.Split(r.RemoteAddr, ":")[0]] = fname
	io.WriteString(w, string(contents))
}

//Potential issue here because POST requests tend to have size limits
func saveFile(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	name := strings.Trim(r.Form.Get("file"), " ./")
	content := r.Form.Get("content")
	ioutil.WriteFile(config["projectdir"]+name, []byte(content), 0744)
}

func hostnameOut(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	io.WriteString(w, string(hostname))
}

//Called as a goroutine to wait for the close command and kill the process.
func watchClose(s *websocket.Conn, p *os.Process) {
	data := make([]byte, 512)
	n, _ := s.Read(data)
	payload := string(data[:n]) //[:n] to cut out padding
	if payload == "close" {
		p.Signal(syscall.SIGINT)
	}
}

//Receives socket connections to /api/socket and creates a PTY to run the process
func socketServer(s *websocket.Conn) {
	data := make([]byte, 512)
	n, _ := s.Read(data)
	com := exec.Command(config["projectdir"] + string(data[:n])) //[:n] to cut out padding
	pt, err := pty.Start(com)
	if err != nil {
		s.Write([]byte("error: " + err.Error()))
		s.Close()
		return
	}
	go watchClose(s, com.Process)
	for {
		out := make([]byte, 1024)
		n, err := pt.Read(out)
		if err != nil {
			if err.Error() == "read /dev/ptmx: input/output error" || err.Error() == "EOF" {
				break
			}
			s.Write([]byte("error: " + err.Error()))
		} else if n > 0 {
			s.Write([]byte("output: " + string(out[:n])))
		}
		time.Sleep(time.Millisecond * 100) //For minimal CPU impact
	}
	s.Close()
}

var changeSockets = make(map[string][]*websocket.Conn)

//Receives socket connections to /api/change and creates a PTY to run the process
func changeServer(s *websocket.Conn) {
	var currFile string
	for {
		var data string
		err := websocket.Message.Receive(s, &data)
		if err != nil {
			s.Close()
			return
		}
		//Change of file
		if data[:3] == "COF:" {
			currFile = data[3:]
			changeSockets[currFile] = append(changeSockets[currFile], s)
		} else if data[:3] == "CIF:" { //Change in file
			for _, v := range changeSockets[currFile] {
				websocket.Message.Send(v, data[:3])
			}
		}
	}
	s.Close()
}
<|endoftext|>


<commit_before>package helixdns

import (
  "github.com/miekg/dns"
  "github.com/coreos/go-etcd/etcd"
  "log"
  "net"
  "strconv"
  "strings"
)

type HelixServer struct {
  Port   int
  Client *etcd.Client
}

func Server(port int, etcdurl string) *HelixServer {
  client := etcd.NewClient([]string{ etcdurl  })
  return &HelixServer {
    Port: port,
    Client: client,
  }
}

func (s HelixServer) Start() {
  handler := newHandler(s.Client)
  server := &dns.Server{
    Addr:         ":"+strconv.Itoa(s.Port),
    Net:          "udp",
    Handler:      dns.HandlerFunc(handler),
    ReadTimeout:  10,
    WriteTimeout: 10,
  }

  log.Print("Starting server...")

  server.ListenAndServe()
}

func getResponse(client *etcd.Client, q dns.Question) (*etcd.Response, error) {
  addr := dns.SplitDomainName(q.Name)
  path := []string{"helix"}

  for s := range addr {
    path = append(path, addr[len(addr)-s-1])
  }

  path = append(path, dns.TypeToString[q.Qtype])

  return client.Get(strings.Join(path, "/"), false, false)
}

func newHandler(client *etcd.Client) func(dns.ResponseWriter, *dns.Msg) {
  return func (w dns.ResponseWriter, req *dns.Msg) {
    m := new(dns.Msg)
    m.SetReply(req)

    qType  := req.Question[0].Qtype
    qClass := req.Question[0].Qclass

    header := dns.RR_Header{Name: m.Question[0].Name, Rrtype: qType, Class: qClass, Ttl: 5}

    resp, err := getResponse(client, req.Question[0])

    if err == nil && qType == dns.TypeA {
      m.Answer = make([]dns.RR, 1)
      m.Answer[0] = &dns.A {Hdr: header, A: net.ParseIP(resp.Node.Value)}
    }

    w.WriteMsg(m)
  }
}
<commit_msg>add AAAA, tweak reversal of labels<commit_after>package helixdns

import (
  "github.com/miekg/dns"
  "github.com/coreos/go-etcd/etcd"
  "log"
  "net"
  "strconv"
  "strings"
)

type HelixServer struct {
  Port   int
  Client *etcd.Client
}

func Server(port int, etcdurl string) *HelixServer {
  client := etcd.NewClient([]string{ etcdurl  })
  return &HelixServer {
    Port: port,
    Client: client,
  }
}

func (s HelixServer) Start() {
  handler := newHandler(s.Client)
  server := &dns.Server{
    Addr:         ":"+strconv.Itoa(s.Port),
    Net:          "udp",
    Handler:      dns.HandlerFunc(handler),
    ReadTimeout:  10,
    WriteTimeout: 10,
  }

  log.Print("Starting server...")

  server.ListenAndServe()
}

func getResponse(client *etcd.Client, q dns.Question) (*etcd.Response, error) {
  addr := dns.SplitDomainName(q.Name)
  path := []string{"helix"}

  for i := len(addr) - 1; i >= 0; i-- {
    path = append(path, addr[i])
  }

  path = append(path, dns.TypeToString[q.Qtype])

  return client.Get(strings.Join(path, "/"), false, false)
}

func newHandler(client *etcd.Client) func(dns.ResponseWriter, *dns.Msg) {
  return func (w dns.ResponseWriter, req *dns.Msg) {
    m := new(dns.Msg)
    m.SetReply(req)

    qType  := req.Question[0].Qtype
    qClass := req.Question[0].Qclass

    header := dns.RR_Header{Name: m.Question[0].Name, Rrtype: qType, Class: qClass, Ttl: 5}

    resp, err := getResponse(client, req.Question[0])

    if err != nil {
      log.Printf("Could not get record for %s", req.Question[0].Name)
      w.WriteMsg(m)
      return
    }

    switch qType {
      case dns.TypeA:
        m.Answer = make([]dns.RR, 1)
        m.Answer[0] = &dns.A {Hdr: header, A: net.ParseIP(resp.Node.Value)}
      case dns.TypeAAAA:
        m.Answer = make([]dns.RR, 1)
        m.Answer[0] = &dns.AAAA {Hdr: header, AAAA: net.ParseIP(resp.Node.Value)}
      default:
        log.Printf("Unrecognised record type: %d",qType)
    }

    w.WriteMsg(m)
  }
}
<|endoftext|>