<commit_before>// Copyright 2014 Jack Wakefield
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"net/http"
	"net/url"

	"github.com/jackwakefield/gopac"
	"github.com/spf13/viper"
	"gopkg.in/elazarl/goproxy.v1"
)

type server struct {
	parser    *gopac.Parser
	client    *http.Client
	http      *goproxy.ProxyHttpServer
	unproxied *http.Client
}

func newServer(parser *gopac.Parser, client *http.Client) (s *server) {
	s = &server{
		parser: parser,
		client: client,
		http:   goproxy.NewProxyHttpServer(),
	}

	s.http.NonproxyHandler = s.nonProxyHandler()

	// determine whether a proxy entry is available for the URL/host for each
	// HTTP request and if so, forward it to the proxy client
	s.http.OnRequest(s.isForwardable()).DoFunc(s.forward)

	return
}

func (s *server) isForwardable() goproxy.ReqConditionFunc {
	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		// find a proxy entry for the request URL/host
		entry, err := s.parser.FindProxy(req.URL.String(), req.URL.Host)

		// ensure no error occurred when finding an entry and that the entry
		// doesn't indicate the connection should be made directly
		if err == nil && entry != "DIRECT" {
			return true
		}

		return false
	}
}

func (s *server) forward(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
	// retrieve a response for the request using the proxy client
	if resp, err := s.client.Get(r.URL.String()); err == nil {
		return r, resp
	}

	return r, nil
}

func (s *server) nonProxyHandler() http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		// set the absolute URL from the host and relative URL and pass it back
		// to the proxy server
		var err error
		req.RequestURI = ""

		if req.URL, err = url.Parse("http://" + req.Host + req.URL.Path); err != nil {
			logger.Errorf("Failed to parse non-proxy URL (%s)", err)
			return
		}

		s.http.ServeHTTP(w, req)
	})
}

func (s *server) address() string {
	return fmt.Sprintf(":%d", viper.GetInt(serverPortKey))
}

func (s *server) listen() error {
	address := s.address()
	logger.Infof("Proxy server listening on address '%s'", address)

	return http.ListenAndServe(address, s.http)
}
<commit_msg>chore(server): set URL host rather than rebuild it<commit_after>// Copyright 2014 Jack Wakefield
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"net/http"

	"github.com/jackwakefield/gopac"
	"github.com/spf13/viper"
	"gopkg.in/elazarl/goproxy.v1"
)

type server struct {
	parser    *gopac.Parser
	client    *http.Client
	http      *goproxy.ProxyHttpServer
	unproxied *http.Client
}

func newServer(parser *gopac.Parser, client *http.Client) (s *server) {
	s = &server{
		parser: parser,
		client: client,
		http:   goproxy.NewProxyHttpServer(),
	}

	s.http.NonproxyHandler = s.nonProxyHandler()

	// determine whether a proxy entry is available for the URL/host for each
	// HTTP request and if so, forward it to the proxy client
	s.http.OnRequest(s.isForwardable()).DoFunc(s.forward)

	return
}

func (s *server) isForwardable() goproxy.ReqConditionFunc {
	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		// find a proxy entry for the request URL/host
		entry, err := s.parser.FindProxy(req.URL.String(), req.URL.Host)

		// ensure no error occurred when finding an entry and that the entry
		// doesn't indicate the connection should be made directly
		if err == nil && entry != "DIRECT" {
			return true
		}

		return false
	}
}

func (s *server) forward(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
	// retrieve a response for the request using the proxy client
	if resp, err := s.client.Get(r.URL.String()); err == nil {
		return r, resp
	}

	return r, nil
}

func (s *server) nonProxyHandler() http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		// set the absolute URL from the host and relative URL and pass it back
		// to the proxy server
		var err error
		req.RequestURI = ""
		req.URL.Host = req.Host

		s.http.ServeHTTP(w, req)
	})
}

func (s *server) address() string {
	return fmt.Sprintf(":%d", viper.GetInt(serverPortKey))
}

func (s *server) listen() error {
	address := s.address()
	logger.Infof("Proxy server listening on address '%s'", address)

	return http.ListenAndServe(address, s.http)
}
<|endoftext|>


<commit_before>package main

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

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

const (
	STACK_LINE = 0
	BODY_START = 2
)

func runServer() {

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

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

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

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

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

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

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

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

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

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

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

		j.stack.condPush(function)

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

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

		}
		j.proceed()
	}
}

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

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

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

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

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

	return filename, line, function
}

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

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

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

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

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

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

	j.files[filename] = lines
}

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

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

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

type stack struct {
	s   []string
	Pos int
}

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

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

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

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

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

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

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

const (
	STACK_LINE = 0
	BODY_START = 2
)

func runServer() {

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

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

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

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

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

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

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

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

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

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

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

		j.stack.condPush(function)

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

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

		}
		j.proceed()
	}
}

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

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

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

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

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

	return filename, line, function
}

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

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

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

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

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

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

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

	j.files[filename] = lines
}

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

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

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

type stack struct {
	s   []string
	Pos int
}

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

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

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

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

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


<commit_before>// Package gatewayrpc is a wrapper around a gorilla/rpc/v2 server which
// registers a special endpoint, "RPC.GetServices", which returns a structure
// containing a description of all services and their methods the rpc server
// supports and their type signatures
package gatewayrpc

import (
	"errors"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/gorilla/rpc/v2"
	"github.com/levenlabs/go-llog"
)

// Server is a simple wrapper around the normal gorilla/rpc/v2 server,
// adding a couple extra features
type Server struct {
	*rpc.Server
	services []Service
}

// NewServer returns a new Server struct initialized with a gorilla/rpc/v2
// server
func NewServer() *Server {
	ns := &Server{Server: rpc.NewServer()}
	ns.Server.RegisterService(ns, "RPC")
	return ns
}

// GetServicesRes describes the structure returned from the GetServices api call
type GetServicesRes struct {
	Services []Service `json:"services"`
}

// GetServices is the actual rpc method which returns the set of services and
// their methods which are supported
func (s *Server) GetServices(r *http.Request, _ *struct{}, res *GetServicesRes) error {
	res.Services = s.services
	return nil
}

// RegisterService passes its arguments through to the underlying gorilla/rpc/v2
// server, as well as adds the given receiver's rpc methods to the Server's
// cache of method data which will be returned by the "RPC.GetMethods" endpoint.
func (s *Server) RegisterService(receiver interface{}, name string) error {
	if err := s.Server.RegisterService(receiver, name); err != nil {
		return err
	}

	name, err := getName(receiver, name)
	if err != nil {
		return err
	}

	service := Service{
		Name:    name,
		Methods: map[string]Method{},
	}
	llog.Debug("retrieving methods")
	for _, method := range getMethods(receiver) {
		llog.Debug("got method", llog.KV{"method": method.Name})
		methodT := method.Type
		args, err := processType(methodT.In(2))
		if err != nil {
			return err
		}
		res, err := processType(methodT.In(3))
		if err != nil {
			return err
		}
		service.Methods[method.Name] = Method{
			Name:    method.Name,
			Args:    args,
			Returns: res,
		}
	}

	s.services = append(s.services, service)

	return nil
}

var (
	typeOfError          = reflect.TypeOf((*error)(nil)).Elem()
	typeOfRequest        = reflect.TypeOf((*http.Request)(nil)).Elem()
	typeOfEmptyInterface = reflect.TypeOf((*interface{})(nil)).Elem()
)

// Since name can optionally be specified to overwrite the name of rcv
func getName(rcv interface{}, name string) (string, error) {
	if name != "" {
		return name, nil
	}
	t := reflect.TypeOf(rcv)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	rcvName := t.Name()
	if !isExported(rcvName) {
		return "", errors.New("receiver not exported")
	}
	return rcvName, nil
}

func getMethods(rcv interface{}) []reflect.Method {
	var ret []reflect.Method
	t := reflect.TypeOf(rcv)
	for i := 0; i < t.NumMethod(); i++ {
		method := t.Method(i)
		mtype := method.Type
		// Method must be exported.
		if method.PkgPath != "" {
			continue
		}
		// Method needs four ins: receiver, *http.Request, *args, *reply.
		if mtype.NumIn() != 4 {
			continue
		}
		// First argument must be a pointer and must be http.Request.
		reqType := mtype.In(1)
		if reqType.Kind() != reflect.Ptr || reqType.Elem() != typeOfRequest {
			continue
		}
		// Second argument must be a pointer and must be exported.
		args := mtype.In(2)
		if args.Kind() != reflect.Ptr || !isExportedOrBuiltin(args) {
			continue
		}
		// Third argument must be a pointer and must be exported.
		reply := mtype.In(3)
		if reply.Kind() != reflect.Ptr || !isExportedOrBuiltin(reply) {
			continue
		}
		// Method needs one out: error.
		if mtype.NumOut() != 1 {
			continue
		}
		if returnType := mtype.Out(0); returnType != typeOfError {
			continue
		}
		ret = append(ret, method)
	}
	return ret
}

func processType(t reflect.Type) (*Type, error) {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	kind := t.Kind()

	// Bool through floats encompasses all integer and float types. Plus string
	if (kind >= reflect.Bool && kind <= reflect.Float64) || kind == reflect.String {
		return &Type{TypeOf: kind}, nil
	}

	if kind == reflect.Array || kind == reflect.Slice {
		innerT, err := processType(t.Elem())
		if err != nil {
			return nil, err
		}
		return &Type{ArrayOf: innerT}, nil
	}

	if kind == reflect.Map {
		if t.Key().Kind() != reflect.String {
			return nil, fmt.Errorf("unsupported map type: %v", t)
		}

		innerT, err := processType(t.Elem())
		if err != nil {
			return nil, err
		}
		return &Type{MapOf: innerT}, nil
	}

	if kind == reflect.Interface {
		if t != typeOfEmptyInterface {
			return nil, fmt.Errorf("unsupported interface: %v", t)
		}
		return &Type{TypeOf: kind}, nil
	}

	if kind == reflect.Struct {
		m := map[string]*Type{}
		for i := 0; i < t.NumField(); i++ {
			f := t.Field(i)
			if !isExported(f.Name) {
				continue
			}
			key := getFieldKey(f)
			innerT, err := processType(f.Type)
			if err != nil {
				return nil, err
			}

			if f.Anonymous && len(innerT.ObjectOf) > 0 {
				for k, v := range innerT.ObjectOf {
					m[k] = v
				}
			} else {
				m[key] = innerT
			}
		}
		return &Type{ObjectOf: m}, nil
	}

	return nil, fmt.Errorf("unsupported type: %v", t)
}

func getFieldKey(f reflect.StructField) string {
	key := f.Name
	jsonTag := f.Tag.Get("json")
	if jsonTag == "" {
		return key
	}

	parts := strings.SplitN(jsonTag, ",", 2)
	if len(parts) == 0 {
		return key
	} else if parts[0] == "" {
		return key
	}

	return parts[0]
}

// isExported returns true of a string is an exported (upper case) name.
func isExported(name string) bool {
	r, _ := utf8.DecodeRuneInString(name)
	return unicode.IsUpper(r)
}

// isExportedOrBuiltin returns true if a type is exported or a builtin.
func isExportedOrBuiltin(t reflect.Type) bool {
	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	// PkgPath will be non-empty even for an exported type,
	// so we need to check the type name as well.
	return isExported(t.Name()) || t.PkgPath() == ""
}
<commit_msg>better errors back, be less strict about processing interfaces<commit_after>// Package gatewayrpc is a wrapper around a gorilla/rpc/v2 server which
// registers a special endpoint, "RPC.GetServices", which returns a structure
// containing a description of all services and their methods the rpc server
// supports and their type signatures
package gatewayrpc

import (
	"errors"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/gorilla/rpc/v2"
	"github.com/levenlabs/go-llog"
)

// Server is a simple wrapper around the normal gorilla/rpc/v2 server,
// adding a couple extra features
type Server struct {
	*rpc.Server
	services []Service
}

// NewServer returns a new Server struct initialized with a gorilla/rpc/v2
// server
func NewServer() *Server {
	ns := &Server{Server: rpc.NewServer()}
	ns.Server.RegisterService(ns, "RPC")
	return ns
}

// GetServicesRes describes the structure returned from the GetServices api call
type GetServicesRes struct {
	Services []Service `json:"services"`
}

// GetServices is the actual rpc method which returns the set of services and
// their methods which are supported
func (s *Server) GetServices(r *http.Request, _ *struct{}, res *GetServicesRes) error {
	res.Services = s.services
	return nil
}

// RegisterService passes its arguments through to the underlying gorilla/rpc/v2
// server, as well as adds the given receiver's rpc methods to the Server's
// cache of method data which will be returned by the "RPC.GetMethods" endpoint.
func (s *Server) RegisterService(receiver interface{}, name string) error {
	if err := s.Server.RegisterService(receiver, name); err != nil {
		return err
	}

	name, err := getName(receiver, name)
	if err != nil {
		return err
	}

	service := Service{
		Name:    name,
		Methods: map[string]Method{},
	}
	llog.Debug("retrieving methods")
	for _, method := range getMethods(receiver) {
		llog.Debug("got method", llog.KV{"method": method.Name})
		methodT := method.Type
		args, err := processType(methodT.In(2))
		if err != nil {
			return fmt.Errorf("processing %q: %s", method.Name, err)
		}
		res, err := processType(methodT.In(3))
		if err != nil {
			return fmt.Errorf("processing %q: %s", method.Name, err)
		}
		service.Methods[method.Name] = Method{
			Name:    method.Name,
			Args:    args,
			Returns: res,
		}
	}

	s.services = append(s.services, service)

	return nil
}

var (
	typeOfError   = reflect.TypeOf((*error)(nil)).Elem()
	typeOfRequest = reflect.TypeOf((*http.Request)(nil)).Elem()
)

// Since name can optionally be specified to overwrite the name of rcv
func getName(rcv interface{}, name string) (string, error) {
	if name != "" {
		return name, nil
	}
	t := reflect.TypeOf(rcv)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	rcvName := t.Name()
	if !isExported(rcvName) {
		return "", errors.New("receiver not exported")
	}
	return rcvName, nil
}

func getMethods(rcv interface{}) []reflect.Method {
	var ret []reflect.Method
	t := reflect.TypeOf(rcv)
	for i := 0; i < t.NumMethod(); i++ {
		method := t.Method(i)
		mtype := method.Type
		// Method must be exported.
		if method.PkgPath != "" {
			continue
		}
		// Method needs four ins: receiver, *http.Request, *args, *reply.
		if mtype.NumIn() != 4 {
			continue
		}
		// First argument must be a pointer and must be http.Request.
		reqType := mtype.In(1)
		if reqType.Kind() != reflect.Ptr || reqType.Elem() != typeOfRequest {
			continue
		}
		// Second argument must be a pointer and must be exported.
		args := mtype.In(2)
		if args.Kind() != reflect.Ptr || !isExportedOrBuiltin(args) {
			continue
		}
		// Third argument must be a pointer and must be exported.
		reply := mtype.In(3)
		if reply.Kind() != reflect.Ptr || !isExportedOrBuiltin(reply) {
			continue
		}
		// Method needs one out: error.
		if mtype.NumOut() != 1 {
			continue
		}
		if returnType := mtype.Out(0); returnType != typeOfError {
			continue
		}
		ret = append(ret, method)
	}
	return ret
}

func processType(t reflect.Type) (*Type, error) {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	kind := t.Kind()

	// Bool through floats encompasses all integer and float types. Plus string
	if (kind >= reflect.Bool && kind <= reflect.Float64) || kind == reflect.String {
		return &Type{TypeOf: kind}, nil
	}

	if kind == reflect.Array || kind == reflect.Slice {
		innerT, err := processType(t.Elem())
		if err != nil {
			return nil, err
		}
		return &Type{ArrayOf: innerT}, nil
	}

	if kind == reflect.Map {
		if t.Key().Kind() != reflect.String {
			return nil, fmt.Errorf("unsupported map type: %v", t)
		}

		innerT, err := processType(t.Elem())
		if err != nil {
			return nil, err
		}
		return &Type{MapOf: innerT}, nil
	}

	if kind == reflect.Interface {
		return &Type{TypeOf: kind}, nil
	}

	if kind == reflect.Struct {
		m := map[string]*Type{}
		for i := 0; i < t.NumField(); i++ {
			f := t.Field(i)
			if !isExported(f.Name) {
				continue
			}
			key := getFieldKey(f)
			innerT, err := processType(f.Type)
			if err != nil {
				return nil, err
			}

			if f.Anonymous && len(innerT.ObjectOf) > 0 {
				for k, v := range innerT.ObjectOf {
					m[k] = v
				}
			} else {
				m[key] = innerT
			}
		}
		return &Type{ObjectOf: m}, nil
	}

	return nil, fmt.Errorf("unsupported type: %v", t)
}

func getFieldKey(f reflect.StructField) string {
	key := f.Name
	jsonTag := f.Tag.Get("json")
	if jsonTag == "" {
		return key
	}

	parts := strings.SplitN(jsonTag, ",", 2)
	if len(parts) == 0 {
		return key
	} else if parts[0] == "" {
		return key
	}

	return parts[0]
}

// isExported returns true of a string is an exported (upper case) name.
func isExported(name string) bool {
	r, _ := utf8.DecodeRuneInString(name)
	return unicode.IsUpper(r)
}

// isExportedOrBuiltin returns true if a type is exported or a builtin.
func isExportedOrBuiltin(t reflect.Type) bool {
	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	// PkgPath will be non-empty even for an exported type,
	// so we need to check the type name as well.
	return isExported(t.Name()) || t.PkgPath() == ""
}
<|endoftext|>


<commit_before>package main

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

	"github.com/byuoitav/authmiddleware"
	"github.com/byuoitav/device-monitoring-microservice/statusinfrastructure"
	"github.com/byuoitav/event-router-microservice/eventinfrastructure"
	"github.com/byuoitav/touchpanel-ui-microservice/events"
	"github.com/byuoitav/touchpanel-ui-microservice/handlers"
	"github.com/byuoitav/touchpanel-ui-microservice/socket"
	"github.com/byuoitav/touchpanel-ui-microservice/uiconfig"
	"github.com/jessemillar/health"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
)

func main() {
	filters := []string{eventinfrastructure.UI, eventinfrastructure.UIFeature}
	en := eventinfrastructure.NewEventNode("Touchpanel UI", filters, os.Getenv("EVENT_ROUTER_ADDRESS"))

	// websocket hub
	hub := socket.NewHub()
	go events.WriteEventsToSocket(en, hub)
	go events.SendRefresh(hub, time.NewTimer(time.Second*10))

	port := ":8888"
	router := echo.New()
	router.Pre(middleware.RemoveTrailingSlash())
	router.Use(middleware.CORS())
	router.Use(echo.WrapMiddleware(authmiddleware.AuthenticateUser))

	router.GET("/health", echo.WrapHandler(http.HandlerFunc(health.Check)))
	router.GET("/mstatus", GetStatus)

	// event endpoints
	router.POST("/publish", handlers.PublishEvent, BindEventNode(en))
	router.POST("/publishfeature", handlers.PublishFeature, BindEventNode(en))

	// websocket
	router.GET("/websocket", func(context echo.Context) error {
		socket.ServeWebsocket(hub, context.Response().Writer, context.Request())
		return nil
	})

	// socket endpoints
	router.PUT("/screenoff", func(context echo.Context) error {
		events.SendScreenTimeout(hub)
		return nil
	})
	router.PUT("/refresh", func(context echo.Context) error {
		events.SendRefresh(hub, time.NewTimer(0))
		return nil
	})
	router.GET("/wsinfo", func(context echo.Context) error {
		si, _ := socket.GetSocketInfo(hub)
		return context.JSON(http.StatusOK, si)
	})
	router.PUT("/socketTest", func(context echo.Context) error {
		events.SendTest(hub)
		return context.JSON(http.StatusOK, "sent")
	})

	router.GET("/pihostname", handlers.GetPiHostname)
	router.GET("/hostname", handlers.GetHostname)
	router.GET("/deviceinfo", handlers.GetDeviceInfo)
	router.GET("/reboot", handlers.Reboot)
	router.GET("/dockerstatus", handlers.GetDockerStatus)

	router.GET("/uiconfig", uiconfig.GetUIConfig)
	router.GET("/uipath", uiconfig.GetUIPath)
	router.GET("/api", uiconfig.GetAPI)
	router.GET("/nextapi", uiconfig.NextAPI)

	router.POST("/help", handlers.Help)
	router.POST("/confirmhelp", handlers.ConfirmHelp)
	router.POST("/cancelhelp", handlers.CancelHelp)

	// all the different ui's
	router.Static("/", "redirect.html")
	router.Any("/404", redirect)
	router.Static("/blueberry", "blueberry-dist")

	router.Start(port)
}

func BindEventNode(en *eventinfrastructure.EventNode) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			c.Set(eventinfrastructure.ContextEventNode, en)
			return next(c)
		}
	}
}

func GetStatus(context echo.Context) error {
	var s statusinfrastructure.Status
	var err error

	s.Version, err = statusinfrastructure.GetVersion("version.txt")
	if err != nil {
		s.Version = "missing"
		s.Status = statusinfrastructure.StatusSick
		s.StatusInfo = fmt.Sprintf("Error: %s", err.Error())
	} else {
		s.Status = statusinfrastructure.StatusOK
		s.StatusInfo = ""
	}

	return context.JSON(http.StatusOK, s)
}

func redirect(context echo.Context) error {
	http.Redirect(context.Response().Writer, context.Request(), "http://github.com/404", 302)
	return nil
}
<commit_msg>remove auth middleware<commit_after>package main

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

	"github.com/byuoitav/device-monitoring-microservice/statusinfrastructure"
	"github.com/byuoitav/event-router-microservice/eventinfrastructure"
	"github.com/byuoitav/touchpanel-ui-microservice/events"
	"github.com/byuoitav/touchpanel-ui-microservice/handlers"
	"github.com/byuoitav/touchpanel-ui-microservice/socket"
	"github.com/byuoitav/touchpanel-ui-microservice/uiconfig"
	"github.com/jessemillar/health"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
)

func main() {
	filters := []string{eventinfrastructure.UI, eventinfrastructure.UIFeature}
	en := eventinfrastructure.NewEventNode("Touchpanel UI", filters, os.Getenv("EVENT_ROUTER_ADDRESS"))

	// websocket hub
	hub := socket.NewHub()
	go events.WriteEventsToSocket(en, hub)
	go events.SendRefresh(hub, time.NewTimer(time.Second*10))

	port := ":8888"
	router := echo.New()
	router.Pre(middleware.RemoveTrailingSlash())
	router.Use(middleware.CORS())
	//	router.Use(echo.WrapMiddleware(authmiddleware.AuthenticateUser))

	router.GET("/health", echo.WrapHandler(http.HandlerFunc(health.Check)))
	router.GET("/mstatus", GetStatus)

	// event endpoints
	router.POST("/publish", handlers.PublishEvent, BindEventNode(en))
	router.POST("/publishfeature", handlers.PublishFeature, BindEventNode(en))

	// websocket
	router.GET("/websocket", func(context echo.Context) error {
		socket.ServeWebsocket(hub, context.Response().Writer, context.Request())
		return nil
	})

	// socket endpoints
	router.PUT("/screenoff", func(context echo.Context) error {
		events.SendScreenTimeout(hub)
		return nil
	})
	router.PUT("/refresh", func(context echo.Context) error {
		events.SendRefresh(hub, time.NewTimer(0))
		return nil
	})
	router.GET("/wsinfo", func(context echo.Context) error {
		si, _ := socket.GetSocketInfo(hub)
		return context.JSON(http.StatusOK, si)
	})
	router.PUT("/socketTest", func(context echo.Context) error {
		events.SendTest(hub)
		return context.JSON(http.StatusOK, "sent")
	})

	router.GET("/pihostname", handlers.GetPiHostname)
	router.GET("/hostname", handlers.GetHostname)
	router.GET("/deviceinfo", handlers.GetDeviceInfo)
	router.GET("/reboot", handlers.Reboot)
	router.GET("/dockerstatus", handlers.GetDockerStatus)

	router.GET("/uiconfig", uiconfig.GetUIConfig)
	router.GET("/uipath", uiconfig.GetUIPath)
	router.GET("/api", uiconfig.GetAPI)
	router.GET("/nextapi", uiconfig.NextAPI)

	router.POST("/help", handlers.Help)
	router.POST("/confirmhelp", handlers.ConfirmHelp)
	router.POST("/cancelhelp", handlers.CancelHelp)

	// all the different ui's
	router.Static("/", "redirect.html")
	router.Any("/404", redirect)
	router.Static("/blueberry", "blueberry-dist")

	router.Start(port)
}

func BindEventNode(en *eventinfrastructure.EventNode) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			c.Set(eventinfrastructure.ContextEventNode, en)
			return next(c)
		}
	}
}

func GetStatus(context echo.Context) error {
	var s statusinfrastructure.Status
	var err error

	s.Version, err = statusinfrastructure.GetVersion("version.txt")
	if err != nil {
		s.Version = "missing"
		s.Status = statusinfrastructure.StatusSick
		s.StatusInfo = fmt.Sprintf("Error: %s", err.Error())
	} else {
		s.Status = statusinfrastructure.StatusOK
		s.StatusInfo = ""
	}

	return context.JSON(http.StatusOK, s)
}

func redirect(context echo.Context) error {
	http.Redirect(context.Response().Writer, context.Request(), "http://github.com/404", 302)
	return nil
}
<|endoftext|>


<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
	abended       bool
	sync          chan bool
	dial          dialer
	pingValue     time.Duration
	pingIndex     int
	pingCount     uint32
	pingWindow    [6]time.Duration
	info          *mongoServerInfo
}

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

type mongoServerInfo struct {
	Master bool
	Mongos bool
	Tags   bson.D
}

var defaultServerInfo mongoServerInfo

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,
		info:         &defaultServerInfo,
	}
	// 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, timeout time.Duration) (socket *mongoSocket, abended bool, err error) {
	for {
		server.Lock()
		abended = server.abended
		if server.closed {
			server.Unlock()
			return nil, abended, errServerClosed
		}
		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]
			info := server.info
			server.Unlock()
			err = socket.InitialAcquire(info, timeout)
			if err != nil {
				continue
			}
		} else {
			server.Unlock()
			socket, err = server.Connect(timeout)
			if err == nil {
				server.Lock()
				// We've waited for the Connect, see if we got
				// closed in the meantime
				if server.closed {
					server.Unlock()
					socket.Release()
					socket.Close()
					return nil, abended, 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(timeout time.Duration) (*mongoSocket, error) {
	server.RLock()
	master := server.info.Master
	dial := server.dial
	server.RUnlock()

	log("Establishing new connection to ", server.Addr, "...")
	var conn net.Conn
	var err error
	if dial == nil {
		// Cannot do this because it lacks timeout support. :-(
		//conn, err = net.DialTCP("tcp", nil, server.tcpaddr)
		conn, err = net.DialTimeout("tcp", server.ResolvedAddr, timeout)
	} 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, timeout), 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) SetInfo(info *mongoServerInfo) {
	server.Lock()
	server.info = info
	server.Unlock()
}

func (server *mongoServer) Info() *mongoServerInfo {
	server.Lock()
	info := server.info
	server.Unlock()
	return info
}

func (server *mongoServer) hasTags(serverTags []bson.D) bool {
NextTagSet:
	for _, tags := range serverTags {
	NextReqTag:
		for _, req := range tags {
			for _, has := range server.info.Tags {
				if req.Name == has.Name {
					if req.Value == has.Value {
						continue NextReqTag
					}
					continue NextTagSet
				}
			}
			continue NextTagSet
		}
		return true
	}
	return false
}

var pingDelay = 5 * 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)
		}
		op := op
		socket, _, err := server.AcquireSocket(0, pingDelay)
		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
}

// BestFit returns the best guess of what would be the most interesting
// server to perform operations on at this point in time.
func (servers *mongoServers) BestFit(serverTags []bson.D) *mongoServer {
	var best *mongoServer
	for _, next := range servers.slice {
		if best == nil {
			best = next
			best.RLock()
			if serverTags != nil && !next.info.Mongos && !best.hasTags(serverTags) {
				best.RUnlock()
				best = nil
			}
			continue
		}
		next.RLock()
		swap := false
		switch {
		case serverTags != nil && !next.info.Mongos && !next.hasTags(serverTags):
			// Must have requested tags.
		case next.info.Master != best.info.Master:
			// Prefer slaves.
			swap = best.info.Master
		case absDuration(next.pingValue-best.pingValue) > 15*time.Millisecond:
			// Prefer nearest server.
			swap = next.pingValue < best.pingValue
		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()
		}
	}
	if best != nil {
		best.RUnlock()
	}
	return best
}

func absDuration(d time.Duration) time.Duration {
	if d < 0 {
		return -d
	}
	return d
}
<commit_msg>Increase the socket timeout used by the pinger.<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
	abended       bool
	sync          chan bool
	dial          dialer
	pingValue     time.Duration
	pingIndex     int
	pingCount     uint32
	pingWindow    [6]time.Duration
	info          *mongoServerInfo
}

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

type mongoServerInfo struct {
	Master bool
	Mongos bool
	Tags   bson.D
}

var defaultServerInfo mongoServerInfo

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,
		info:         &defaultServerInfo,
	}
	// 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, timeout time.Duration) (socket *mongoSocket, abended bool, err error) {
	for {
		server.Lock()
		abended = server.abended
		if server.closed {
			server.Unlock()
			return nil, abended, errServerClosed
		}
		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]
			info := server.info
			server.Unlock()
			err = socket.InitialAcquire(info, timeout)
			if err != nil {
				continue
			}
		} else {
			server.Unlock()
			socket, err = server.Connect(timeout)
			if err == nil {
				server.Lock()
				// We've waited for the Connect, see if we got
				// closed in the meantime
				if server.closed {
					server.Unlock()
					socket.Release()
					socket.Close()
					return nil, abended, 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(timeout time.Duration) (*mongoSocket, error) {
	server.RLock()
	master := server.info.Master
	dial := server.dial
	server.RUnlock()

	log("Establishing new connection to ", server.Addr, "...")
	var conn net.Conn
	var err error
	if dial == nil {
		// Cannot do this because it lacks timeout support. :-(
		//conn, err = net.DialTCP("tcp", nil, server.tcpaddr)
		conn, err = net.DialTimeout("tcp", server.ResolvedAddr, timeout)
	} 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, timeout), 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) SetInfo(info *mongoServerInfo) {
	server.Lock()
	server.info = info
	server.Unlock()
}

func (server *mongoServer) Info() *mongoServerInfo {
	server.Lock()
	info := server.info
	server.Unlock()
	return info
}

func (server *mongoServer) hasTags(serverTags []bson.D) bool {
NextTagSet:
	for _, tags := range serverTags {
	NextReqTag:
		for _, req := range tags {
			for _, has := range server.info.Tags {
				if req.Name == has.Name {
					if req.Value == has.Value {
						continue NextReqTag
					}
					continue NextTagSet
				}
			}
			continue NextTagSet
		}
		return true
	}
	return false
}

var pingDelay = 5 * 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)
		}
		op := op
		socket, _, err := server.AcquireSocket(0, 3 * pingDelay)
		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
}

// BestFit returns the best guess of what would be the most interesting
// server to perform operations on at this point in time.
func (servers *mongoServers) BestFit(serverTags []bson.D) *mongoServer {
	var best *mongoServer
	for _, next := range servers.slice {
		if best == nil {
			best = next
			best.RLock()
			if serverTags != nil && !next.info.Mongos && !best.hasTags(serverTags) {
				best.RUnlock()
				best = nil
			}
			continue
		}
		next.RLock()
		swap := false
		switch {
		case serverTags != nil && !next.info.Mongos && !next.hasTags(serverTags):
			// Must have requested tags.
		case next.info.Master != best.info.Master:
			// Prefer slaves.
			swap = best.info.Master
		case absDuration(next.pingValue-best.pingValue) > 15*time.Millisecond:
			// Prefer nearest server.
			swap = next.pingValue < best.pingValue
		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()
		}
	}
	if best != nil {
		best.RUnlock()
	}
	return best
}

func absDuration(d time.Duration) time.Duration {
	if d < 0 {
		return -d
	}
	return d
}
<|endoftext|>


<commit_before>package nodeflock

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"strconv"
	"time"
)

type server struct {
	flock *Flock
}

func (s *server) handleConnection(conn net.Conn) error {
	var headerBuffer [1]byte
	header := ""
	for {
		_, readErr := conn.Read(headerBuffer[0:])
		if readErr != nil {
			return readErr
		}
		// let us read with 1 byte steps on conn until we find "{"
		if headerBuffer[0:1][0] == 123 {
			jsonLength, strconvErr := strconv.Atoi(string(header))
			if strconvErr != nil {
				return strconvErr
			}
			bytesRead := 1
			jsonBytes := []byte{123}
			for bytesRead < jsonLength {
				nextJSONBytes := make([]byte, 1024)
				n, readErr := conn.Read(nextJSONBytes)
				if readErr != nil && readErr != io.EOF {
					return readErr
				}
				if readErr == io.EOF && bytesRead < jsonLength {
					return errors.New(fmt.Sprintln("eof, but not done, expected", jsonLength, "from", header, "got", bytesRead, "::", string(jsonBytes)))
				}
				jsonBytes = append(jsonBytes, nextJSONBytes[0:n]...)
				bytesRead += n

			}
			//fmt.Println("raw call", string(jsonBytes), jsonBytes)
			resultBytes, errRawCall := s.flock.CallRaw(jsonBytes)
			if errRawCall != nil {
				return errRawCall
			}
			resultBytes = append([]byte(fmt.Sprint(len(resultBytes))), resultBytes...)
			//fmt.Println("ready to reply", string(resultBytes))
			bytesWritten := 0
			for bytesWritten < len(resultBytes) {
				// respond
				n, writeErr := conn.Write(resultBytes[bytesWritten:])
				if writeErr != nil {
					return writeErr
				}
				bytesWritten += n
			}
			//fmt.Println("done here")
			header = ""
			continue
		}
		header += string(headerBuffer[0:])
	}
}

func RunServer(addr string, flockSize int, jsModuleFile string, maxExecutionTime time.Duration) error {
	flock, flockErr := NewFlock(jsModuleFile, flockSize, maxExecutionTime)
	if flockErr != nil {
		return flockErr
	}
	s := &server{
		flock: flock,
	}
	return s.run(addr, flockSize, jsModuleFile)
}

func (s *server) run(addr string, flockSize int, jsModuleFile string) error {

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	for {
		// this blocks until connection or error
		conn, err := ln.Accept()
		if err != nil {
			log.Println("RunSocketServer: could not accept connection", fmt.Sprint(err))
			continue
		}
		// a goroutine handles conn so that the loop can accept other connections
		go func() {
			handlingErr := s.handleConnection(conn)
			if handlingErr != nil && handlingErr != io.EOF {
				fmt.Println("could not handle request in socket server::", handlingErr)
			}
			closeErr := conn.Close()
			if closeErr != nil {
				fmt.Println("could not close connection", closeErr)
			}
		}()
	}
}
<commit_msg>safer connection handling<commit_after>package nodeflock

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"strconv"
	"time"
)

type server struct {
	flock *Flock
}

func (s *server) handleConnection(conn net.Conn) error {
	var headerBuffer [1]byte
	header := ""
	for {
		_, readErr := conn.Read(headerBuffer[0:])
		if readErr != nil {
			return readErr
		}
		// let us read with 1 byte steps on conn until we find "{"
		if headerBuffer[0:1][0] == 123 {
			jsonLength, strconvErr := strconv.Atoi(string(header))
			if strconvErr != nil {
				return strconvErr
			}
			bytesRead := 1
			jsonBytes := []byte{123}
			for bytesRead < jsonLength {
				nextJSONBytes := make([]byte, 1024)
				n, readErr := conn.Read(nextJSONBytes)
				if readErr != nil && readErr != io.EOF {
					return readErr
				}
				if readErr == io.EOF && bytesRead < jsonLength {
					return errors.New(fmt.Sprintln("eof, but not done, expected", jsonLength, "from", header, "got", bytesRead, "::", string(jsonBytes)))
				}
				jsonBytes = append(jsonBytes, nextJSONBytes[0:n]...)
				bytesRead += n

			}
			//fmt.Println("raw call", string(jsonBytes), jsonBytes)
			resultBytes, errRawCall := s.flock.CallRaw(jsonBytes)
			if errRawCall != nil {
				return errRawCall
			}
			resultBytes = append([]byte(fmt.Sprint(len(resultBytes))), resultBytes...)
			//fmt.Println("ready to reply", string(resultBytes))
			bytesWritten := 0
			for bytesWritten < len(resultBytes) {
				// respond
				n, writeErr := conn.Write(resultBytes[bytesWritten:])
				if writeErr != nil {
					return writeErr
				}
				bytesWritten += n
			}
			//fmt.Println("done here")
			header = ""
			continue
		}
		header += string(headerBuffer[0:])
	}
}

func RunServer(addr string, flockSize int, jsModuleFile string, maxExecutionTime time.Duration) error {
	flock, flockErr := NewFlock(jsModuleFile, flockSize, maxExecutionTime)
	if flockErr != nil {
		return flockErr
	}
	s := &server{
		flock: flock,
	}
	return s.run(addr, flockSize, jsModuleFile)
}

func (s *server) run(addr string, flockSize int, jsModuleFile string) error {

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	for {
		// this blocks until connection or error
		nextConn, err := ln.Accept()
		if err != nil {
			log.Println("RunSocketServer: could not accept connection", fmt.Sprint(err))
			continue
		}

		// a goroutine handles conn so that the loop can accept other connections
		go func(conn net.Conn) {

			handlingErr := s.handleConnection(conn)
			if handlingErr != nil && handlingErr != io.EOF {
				fmt.Println("could not handle request in socket server::", handlingErr)
			}
			closeErr := conn.Close()
			if closeErr != nil {
				fmt.Println("could not close connection", closeErr)
			}
		}(nextConn)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Eduardo Pinheiro (edpin@edpin.com). All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package https

import (
	"crypto/tls"
	"golang.org/x/crypto/acme/autocert"
	"log"
	"net/http"
	"time"
)

// StartSecureServer starts an HTTPS server with a mux and an autocert manager.
// The HTTPS server started enables HTST by default to ensure maximum protection
// (see https://www.owasp.org/index.php/HTTP_Strict_Transport_Security_Cheat_Sheet).
// StartSecureServer also starts an HTTP server that redirects all requests to
// their HTTPS counterpart and immediately terminates all connections.
func StartSecureServer(mux *http.ServeMux, m *autocert.Manager) {
	s := NewSecureServer(m)
	s.Handler = NewHSTS(mux)
	go func() {
		// Redirect regular HTTP requests to HTTPS.
		insecure := &http.Server{
			ReadTimeout:  5 * time.Second,
			WriteTimeout: 5 * time.Second,
			Handler: http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
				w.Header().Set("Connection", "close")
				url := "https://" + req.Host + req.URL.String()
				http.Redirect(w, req, url, http.StatusMovedPermanently)
			}),
		}
		log.Fatal(insecure.ListenAndServe())
	}()
	log.Fatal(s.ListenAndServeTLS("", ""))
}

type htstMux struct {
	*http.ServeMux
}

// NewSecureServer returns a new HTTP server with strict security settings.
func NewSecureServer(m *autocert.Manager) *http.Server {
	t := m.TLSConfig()
	t.ClientSessionCache = tls.NewLRUClientSessionCache(0)
	t.MinVersion = tls.VersionTLS12
	t.CurvePreferences = []tls.CurveID{
		tls.X25519, // requires go 1.8
		tls.CurveP521,
		tls.CurveP384,
		tls.CurveP256,
	}
	// Prefer this order of ciphers.
	t.CipherSuites = []uint16{
		tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
		tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305,
		// required by HTTP-2.
		tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
	}

	return &http.Server{
		Addr:              ":https",
		ReadTimeout:       5 * time.Second,
		ReadHeaderTimeout: 1 * time.Second,
		WriteTimeout:      10 * time.Second,
		IdleTimeout:       120 * time.Second,
		TLSConfig:         t,
	}
}

// NewHSTS returns an HTTP handler that sets HSTS headers on all requests.
func NewHSTS(mux *http.ServeMux) http.Handler {
	return htstMux{
		ServeMux: mux,
	}
}

// ServeHTTP implements http.Handler.
func (h htstMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Strict-Transport-Security", "max-age=31536000; includeSubDomains")
	h.ServeMux.ServeHTTP(w, r)
}
<commit_msg>uses the less restrictive interface (Handler) instead of a ServeMux<commit_after>// Copyright 2017 Eduardo Pinheiro (edpin@edpin.com). All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package https

import (
	"crypto/tls"
	"golang.org/x/crypto/acme/autocert"
	"log"
	"net/http"
	"time"
)

// StartSecureServer starts an HTTPS server with a Handler and an autocert
// manager. The HTTPS server started enables HTST by default to ensure maximum
// protection (see
// https://www.owasp.org/index.php/HTTP_Strict_Transport_Security_Cheat_Sheet).
// StartSecureServer also starts an HTTP server that redirects all requests to
// their HTTPS counterpart and immediately terminates all connections.
func StartSecureServer(h http.Handler, m *autocert.Manager) {
	s := NewSecureServer(m)
	s.Handler = NewHSTS(h)
	go func() {
		// Redirect regular HTTP requests to HTTPS.
		insecure := &http.Server{
			ReadTimeout:  5 * time.Second,
			WriteTimeout: 5 * time.Second,
			Handler: http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
				w.Header().Set("Connection", "close")
				url := "https://" + req.Host + req.URL.String()
				http.Redirect(w, req, url, http.StatusMovedPermanently)
			}),
		}
		log.Fatal(insecure.ListenAndServe())
	}()
	log.Fatal(s.ListenAndServeTLS("", ""))
}

type hstsHandler struct {
	http.Handler
}

// NewSecureServer returns a new HTTP server with strict security settings.
func NewSecureServer(m *autocert.Manager) *http.Server {
	t := m.TLSConfig()
	t.ClientSessionCache = tls.NewLRUClientSessionCache(0)
	t.MinVersion = tls.VersionTLS12
	t.CurvePreferences = []tls.CurveID{
		tls.X25519, // requires go 1.8
		tls.CurveP521,
		tls.CurveP384,
		tls.CurveP256,
	}
	// Prefer this order of ciphers.
	t.CipherSuites = []uint16{
		tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
		tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305,
		// required by HTTP-2.
		tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
	}

	return &http.Server{
		Addr:              ":https",
		ReadTimeout:       5 * time.Second,
		ReadHeaderTimeout: 1 * time.Second,
		WriteTimeout:      10 * time.Second,
		IdleTimeout:       120 * time.Second,
		TLSConfig:         t,
	}
}

// NewHSTS returns an HTTP handler that sets HSTS headers on all requests.
func NewHSTS(h http.Handler) http.Handler {
	return hstsHandler{
		h,
	}
}

// ServeHTTP implements http.Handler.
func (h hstsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Strict-Transport-Security", "max-age=31536000; includeSubDomains")
	h.ServeHTTP(w, r)
}
<|endoftext|>


<commit_before>// Package server offers a simple server with logging and configs
package server

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"time"
)

// TODO: Graceful restart would be nice - http://grisha.org/blog/2014/06/03/graceful-restart-in-golang/

// Logger interface for a simple logger (the stdlib log pkg and the fragmenta log pkg conform)
type Logger interface {
	Printf(format string, args ...interface{})
}

// Server holds the config and logger for the app
type Server struct {
	// Our internal logger instance
	Logger Logger

	// Which port to serve on
	port int

	// Which env mode we're in, read from ENV variable
	production bool

	// Server configs - access with Config(string)
	configProduction  map[string]string
	configDevelopment map[string]string
	configTest        map[string]string
}

// New creates a new server instance
func New() (*Server, error) {

	// Check environment variable to see if we are in production mode
	prod := false
	if os.Getenv("FRAG_ENV") == "production" {
		prod = true
	}

	// Set up a new server
	s := &Server{
		port:              3000,
		production:        prod,
		configProduction:  make(map[string]string),
		configDevelopment: make(map[string]string),
		configTest:        make(map[string]string),
		Logger:            log.New(os.Stderr, "fragmenta: ", log.LstdFlags), // default to a stderr logger
	}

	err := s.readConfig()
	if err != nil {
		return s, err
	}
	err = s.readArguments()
	if err != nil {
		return s, err
	}

	return s, err
}

// Logf logs the message with the given arguments to our internal logger
func (s *Server) Logf(format string, v ...interface{}) {
	s.Logger.Printf(format, v...)
}

// Log logs the message to our internal logger
func (s *Server) Log(message string) {
	s.Logf(message)
}

// Fatalf the message with the given arguments to our internal logger, and then exits with status 1
func (s *Server) Fatalf(format string, v ...interface{}) {
	s.Logger.Printf(format, v...)

	// Now exit
	os.Exit(1)
}

// Fatal logs the message, and then exits with status 1
func (s *Server) Fatal(format string) {
	s.Fatalf(format)
}

// Timef logs a time since starting, when used with defer at the start of a function to time
// Usage: defer s.Timef("Completed %s in %s",time.Now(),args...)
func (s *Server) Timef(format string, start time.Time, v ...interface{}) {
	end := time.Since(start).String()
	var args []interface{}
	args = append(args, end)
	args = append(args, v...)
	s.Logf(format, args...)
}

// Start starts an http server on the given port
func (s *Server) Start() error {
	p := fmt.Sprintf(":%d", s.port)
	return http.ListenAndServe(p, nil)
}

// StartTLS starts an https server on the given port
func (s *Server) StartTLS() error {
	port := fmt.Sprintf(":%d", s.port)
	cert := s.Config("tls_cert")
	key := s.Config("tls_key")
	return http.ListenAndServeTLS(port, cert, key, nil)
}
<commit_msg>Added autocert support for let's encrypt certs using StartTLSAutocert, adjusted default server config<commit_after>// Package server offers a simple server with logging and configs
package server

import (
	"crypto/tls"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"golang.org/x/crypto/acme/autocert"
)

// Logger interface for a simple logger (the stdlib log pkg and the fragmenta log pkg conform)
type Logger interface {
	Printf(format string, args ...interface{})
}

// Server holds the config and logger for the app
type Server struct {
	// Our internal logger instance
	Logger Logger

	// Which port to serve on
	port int

	// Which env mode we're in, read from ENV variable
	production bool

	// Server configs - access with Config(string)
	configProduction  map[string]string
	configDevelopment map[string]string
	configTest        map[string]string
}

// New creates a new server instance
func New() (*Server, error) {

	// Check environment variable to see if we are in production mode
	prod := false
	if os.Getenv("FRAG_ENV") == "production" {
		prod = true
	}

	// Set up a new server
	s := &Server{
		port:              3000,
		production:        prod,
		configProduction:  make(map[string]string),
		configDevelopment: make(map[string]string),
		configTest:        make(map[string]string),
		Logger:            log.New(os.Stderr, "fragmenta: ", log.LstdFlags),
	}

	err := s.readConfig()
	if err != nil {
		return s, err
	}
	err = s.readArguments()
	if err != nil {
		return s, err
	}

	return s, err
}

// Logf logs the message with the given arguments to our internal logger
func (s *Server) Logf(format string, v ...interface{}) {
	s.Logger.Printf(format, v...)
}

// Log logs the message to our internal logger
func (s *Server) Log(message string) {
	s.Logf(message)
}

// Fatalf the message with the given arguments to our internal logger, and then exits with status 1
func (s *Server) Fatalf(format string, v ...interface{}) {
	s.Logger.Printf(format, v...)

	// Now exit
	os.Exit(1)
}

// Fatal logs the message, and then exits with status 1
func (s *Server) Fatal(format string) {
	s.Fatalf(format)
}

// Timef logs a time since starting, when used with defer at the start of a function to time
// Usage: defer s.Timef("Completed %s in %s",time.Now(),args...)
func (s *Server) Timef(format string, start time.Time, v ...interface{}) {
	end := time.Since(start).String()
	var args []interface{}
	args = append(args, end)
	args = append(args, v...)
	s.Logf(format, args...)
}

// PortString returns a string port suitable for passing to http.Server
func (s *Server) PortString() string {
	return fmt.Sprintf(":%d", s.port)
}

// Start starts an http server on the given port
func (s *Server) Start() error {
	server := &http.Server{
		Addr:         s.PortString(),
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 10 * time.Second,
	}
	return server.ListenAndServe()
}

// StartTLS starts an https server on the given port
// with tls cert/key from config keys.
func (s *Server) StartTLS() error {
	cert := s.Config("tls_cert")
	key := s.Config("tls_key")
	return s.ConfiguredTLSServer(nil).ListenAndServeTLS(cert, key)
}

// StartTLSAutocert starts an https server on the given port
// by requesting certs from an ACME provider.
// The server must be on a public IP which matches the
// DNS for the domains.
func (s *Server) StartTLSAutocert(email string, domains string) error {
	autocertDomains := strings.Split(domains, " ")
	certManager := &autocert.Manager{
		Prompt:     autocert.AcceptTOS,
		Email:      email,                                      // Email for problems with certs
		HostPolicy: autocert.HostWhitelist(autocertDomains...), // Domains to request certs for
		Cache:      autocert.DirCache("secrets"),               // Cache certs in secrets folder
	}
	server := s.ConfiguredTLSServer(certManager)
	return server.ListenAndServeTLS("", "")
}

// ConfiguredTLSServer returns a TLS server instance with a secure config
// this server has read/write timeouts set to 20 seconds,
// prefers server cipher suites and only uses certain accelerated curves
// see - https://blog.gopheracademy.com/advent-2016/exposing-go-on-the-internet/
func (s *Server) ConfiguredTLSServer(certManager *autocert.Manager) *http.Server {

	return &http.Server{
		// Set the port in the preferred string format
		Addr: s.PortString(),
		/*
			// The default server from net/http has no timeouts
			ReadTimeout:  20 * time.Second,
			WriteTimeout: 20 * time.Second,
			IdleTimeout:  120 * time.Second,
		*/
		// This TLS config follows recommendations in the above article
		TLSConfig: &tls.Config{
			// Pass in a cert manager if you want one set
			// this will only be used if the server Certificates are empty
			GetCertificate: certManager.GetCertificate,

			// VersionTLS11 or VersionTLS12 would exclude many browsers
			// inc. Android 4.x, IE 10, Opera 12.17, Safari 6
			// So unfortunately not acceptable as a default yet
			// Current default here for clarity
			MinVersion: tls.VersionTLS10,

			// Causes servers to use Go's default ciphersuite preferences,
			// which are tuned to avoid attacks. Does nothing on clients.
			PreferServerCipherSuites: true,
			// Only use curves which have assembly implementations
			CurvePreferences: []tls.CurveID{
				tls.CurveP256,
				tls.X25519, // Go 1.8 only
			},
		},
	}

}

// StartRedirectAll starts redirecting from port given to the given url
// this should be called before StartTLS if redirecting to https
func (s *Server) StartRedirectAll(p int, url string) {
	port := fmt.Sprintf(":%d", p)
	// Listen and server on port p in a separate goroutine
	go func() {
		http.ListenAndServe(port, &redirectHandler{redirect: url})
	}()
}

// redirectHandler is useful if serving tls direct (not behind a proxy)
// and a redirect from port 80 is required.
type redirectHandler struct {
	redirect string
}

// ServeHTTP on this handler simply redirects to the main site
func (m *redirectHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	http.Redirect(w, r, m.redirect, http.StatusMovedPermanently)
}
<|endoftext|>


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

// 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()
				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,
	}
	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)
		}
		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>Send pings with op.limit = -1; mongos blows up without it.<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")

// 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()
				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)
		}
		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|>