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

package walk

import (
	"unsafe"
)

import (
	"github.com/lxn/win"
)

type Margins struct {
	HNear, VNear, HFar, VFar int
}

type Layout interface {
	Container() Container
	SetContainer(value Container)
	Margins() Margins
	SetMargins(value Margins) error
	Spacing() int
	SetSpacing(value int) error
	LayoutFlags() LayoutFlags
	MinSize() Size
	Update(reset bool) error
}

func shouldLayoutWidget(widget Widget) bool {
	if widget == nil {
		return false
	}

	_, isSpacer := widget.(*Spacer)

	return isSpacer || widget.AsWindowBase().visible || widget.AlwaysConsumeSpace()
}

func DescendantByName(container Container, name string) Widget {
	var widget Widget

	walkDescendants(container.AsContainerBase(), func(w Window) bool {
		if w.Name() == name {
			widget = w.(Widget)
			return false
		}

		return true
	})

	if widget == nil {
		return nil
	}

	return widget
}

type Container interface {
	Window
	AsContainerBase() *ContainerBase
	Children() *WidgetList
	Layout() Layout
	SetLayout(value Layout) error
	DataBinder() *DataBinder
	SetDataBinder(dbm *DataBinder)
}

type ContainerBase struct {
	WidgetBase
	layout     Layout
	children   *WidgetList
	dataBinder *DataBinder
	persistent bool
}

func (cb *ContainerBase) AsWidgetBase() *WidgetBase {
	return &cb.WidgetBase
}

func (cb *ContainerBase) AsContainerBase() *ContainerBase {
	return cb
}

func (cb *ContainerBase) LayoutFlags() LayoutFlags {
	if cb.layout == nil {
		return 0
	}

	return cb.layout.LayoutFlags()
}

func (cb *ContainerBase) MinSizeHint() Size {
	if cb.layout == nil {
		return Size{}
	}

	return cb.layout.MinSize()
}

func (cb *ContainerBase) SizeHint() Size {
	return Size{100, 100}
}

func (cb *ContainerBase) applyEnabled(enabled bool) {
	cb.WidgetBase.applyEnabled(enabled)

	applyEnabledToDescendants(cb.window.(Widget), enabled)
}

func (cb *ContainerBase) applyFont(font *Font) {
	cb.WidgetBase.applyFont(font)

	applyFontToDescendants(cb.window.(Widget), font)
}

func (cb *ContainerBase) Children() *WidgetList {
	return cb.children
}

func (cb *ContainerBase) Layout() Layout {
	return cb.layout
}

func (cb *ContainerBase) SetLayout(value Layout) error {
	if cb.layout != value {
		if cb.layout != nil {
			cb.layout.SetContainer(nil)
		}

		cb.layout = value

		if value != nil && value.Container() != Container(cb) {
			value.SetContainer(cb)
		}
	}

	return nil
}

func (cb *ContainerBase) DataBinder() *DataBinder {
	return cb.dataBinder
}

func (cb *ContainerBase) SetDataBinder(db *DataBinder) {
	if db == cb.dataBinder {
		return
	}

	if cb.dataBinder != nil {
		cb.dataBinder.SetBoundWidgets(nil)
	}

	cb.dataBinder = db

	if db != nil {
		var boundWidgets []Widget

		walkDescendants(cb.window, func(w Window) bool {
			if w.Handle() == cb.hWnd {
				return true
			}

			if c, ok := w.(Container); ok && c.DataBinder() != nil {
				return false
			}

			for _, prop := range w.AsWindowBase().name2Property {
				if _, ok := prop.Source().(string); ok {
					boundWidgets = append(boundWidgets, w.(Widget))
					break
				}
			}

			return true
		})

		db.SetBoundWidgets(boundWidgets)
	}
}

func (cb *ContainerBase) forEachPersistableChild(f func(p Persistable) error) error {
	if cb.children == nil {
		return nil
	}

	for _, child := range cb.children.items {
		if persistable, ok := child.(Persistable); ok && persistable.Persistent() {
			if err := f(persistable); err != nil {
				return err
			}
		}
	}

	return nil
}

func (cb *ContainerBase) Persistent() bool {
	return cb.persistent
}

func (cb *ContainerBase) SetPersistent(value bool) {
	cb.persistent = value
}

func (cb *ContainerBase) SaveState() error {
	return cb.forEachPersistableChild(func(p Persistable) error {
		return p.SaveState()
	})
}

func (cb *ContainerBase) RestoreState() error {
	return cb.forEachPersistableChild(func(p Persistable) error {
		return p.RestoreState()
	})
}

func (cb *ContainerBase) SetSuspended(suspend bool) {
	wasSuspended := cb.Suspended()

	cb.WidgetBase.SetSuspended(suspend)

	if !suspend && wasSuspended && cb.layout != nil {
		cb.layout.Update(false)
	}
}

func (cb *ContainerBase) WndProc(hwnd win.HWND, msg uint32, wParam, lParam uintptr) uintptr {
	switch msg {
	case win.WM_COMMAND:
		if lParam == 0 {
			switch win.HIWORD(uint32(wParam)) {
			case 0:
				cmdId := win.LOWORD(uint32(wParam))
				switch cmdId {
				case win.IDOK, win.IDCANCEL:
					form := ancestor(cb)
					if form == nil {
						break
					}

					dlg, ok := form.(dialogish)
					if !ok {
						break
					}

					var button *PushButton
					if cmdId == win.IDOK {
						button = dlg.DefaultButton()
					} else {
						button = dlg.CancelButton()
					}

					if button != nil && button.Visible() && button.Enabled() {
						button.raiseClicked()
					}

					break
				}

				// Menu
				actionId := uint16(win.LOWORD(uint32(wParam)))
				if action, ok := actionsById[actionId]; ok {
					action.raiseTriggered()
					return 0
				}

			case 1:
				// Accelerator
			}
		} else {
			// The window that sent the notification shall handle it itself.
			hWnd := win.HWND(lParam)
			if window := windowFromHandle(hWnd); window != nil {
				window.WndProc(hwnd, msg, wParam, lParam)
				return 0
			}
		}

	case win.WM_NOTIFY:
		nmh := (*win.NMHDR)(unsafe.Pointer(lParam))
		if window := windowFromHandle(nmh.HwndFrom); window != nil {
			// The window that sent the notification shall handle it itself.
			return window.WndProc(hwnd, msg, wParam, lParam)
		}

	case win.WM_SIZE, win.WM_SIZING:
		if cb.layout != nil {
			cb.layout.Update(false)
		}
	}

	return cb.WidgetBase.WndProc(hwnd, msg, wParam, lParam)
}

func (cb *ContainerBase) onInsertingWidget(index int, widget Widget) (err error) {
	return nil
}

func (cb *ContainerBase) onInsertedWidget(index int, widget Widget) (err error) {
	if parent := widget.Parent(); parent == nil || parent.Handle() != cb.hWnd {
		if err = widget.SetParent(cb.window.(Container)); err != nil {
			return
		}
	}

	if cb.layout != nil {
		cb.layout.Update(true)
	}

	widget.(applyFonter).applyFont(cb.Font())

	return
}

func (cb *ContainerBase) onRemovingWidget(index int, widget Widget) (err error) {
	if widget.Parent() == nil {
		return
	}

	if widget.Parent().Handle() == cb.hWnd {
		err = widget.SetParent(nil)
	}

	return
}

func (cb *ContainerBase) onRemovedWidget(index int, widget Widget) (err error) {
	if cb.layout != nil {
		cb.layout.Update(true)
	}

	return
}

func (cb *ContainerBase) onClearingWidgets() (err error) {
	for _, widget := range cb.children.items {
		if widget.Parent().Handle() == cb.hWnd {
			if err = widget.SetParent(nil); err != nil {
				return
			}
		}
	}

	return
}

func (cb *ContainerBase) onClearedWidgets() (err error) {
	if cb.layout != nil {
		cb.layout.Update(true)
	}

	return
}
<commit_msg>ContainerBase: Use default SizeHint implementation, which calls MinSizeHint<commit_after>// Copyright 2010 The Walk Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package walk

import (
	"unsafe"
)

import (
	"github.com/lxn/win"
)

type Margins struct {
	HNear, VNear, HFar, VFar int
}

type Layout interface {
	Container() Container
	SetContainer(value Container)
	Margins() Margins
	SetMargins(value Margins) error
	Spacing() int
	SetSpacing(value int) error
	LayoutFlags() LayoutFlags
	MinSize() Size
	Update(reset bool) error
}

func shouldLayoutWidget(widget Widget) bool {
	if widget == nil {
		return false
	}

	_, isSpacer := widget.(*Spacer)

	return isSpacer || widget.AsWindowBase().visible || widget.AlwaysConsumeSpace()
}

func DescendantByName(container Container, name string) Widget {
	var widget Widget

	walkDescendants(container.AsContainerBase(), func(w Window) bool {
		if w.Name() == name {
			widget = w.(Widget)
			return false
		}

		return true
	})

	if widget == nil {
		return nil
	}

	return widget
}

type Container interface {
	Window
	AsContainerBase() *ContainerBase
	Children() *WidgetList
	Layout() Layout
	SetLayout(value Layout) error
	DataBinder() *DataBinder
	SetDataBinder(dbm *DataBinder)
}

type ContainerBase struct {
	WidgetBase
	layout     Layout
	children   *WidgetList
	dataBinder *DataBinder
	persistent bool
}

func (cb *ContainerBase) AsWidgetBase() *WidgetBase {
	return &cb.WidgetBase
}

func (cb *ContainerBase) AsContainerBase() *ContainerBase {
	return cb
}

func (cb *ContainerBase) LayoutFlags() LayoutFlags {
	if cb.layout == nil {
		return 0
	}

	return cb.layout.LayoutFlags()
}

func (cb *ContainerBase) MinSizeHint() Size {
	if cb.layout == nil {
		return Size{}
	}

	return cb.layout.MinSize()
}

func (cb *ContainerBase) applyEnabled(enabled bool) {
	cb.WidgetBase.applyEnabled(enabled)

	applyEnabledToDescendants(cb.window.(Widget), enabled)
}

func (cb *ContainerBase) applyFont(font *Font) {
	cb.WidgetBase.applyFont(font)

	applyFontToDescendants(cb.window.(Widget), font)
}

func (cb *ContainerBase) Children() *WidgetList {
	return cb.children
}

func (cb *ContainerBase) Layout() Layout {
	return cb.layout
}

func (cb *ContainerBase) SetLayout(value Layout) error {
	if cb.layout != value {
		if cb.layout != nil {
			cb.layout.SetContainer(nil)
		}

		cb.layout = value

		if value != nil && value.Container() != Container(cb) {
			value.SetContainer(cb)
		}
	}

	return nil
}

func (cb *ContainerBase) DataBinder() *DataBinder {
	return cb.dataBinder
}

func (cb *ContainerBase) SetDataBinder(db *DataBinder) {
	if db == cb.dataBinder {
		return
	}

	if cb.dataBinder != nil {
		cb.dataBinder.SetBoundWidgets(nil)
	}

	cb.dataBinder = db

	if db != nil {
		var boundWidgets []Widget

		walkDescendants(cb.window, func(w Window) bool {
			if w.Handle() == cb.hWnd {
				return true
			}

			if c, ok := w.(Container); ok && c.DataBinder() != nil {
				return false
			}

			for _, prop := range w.AsWindowBase().name2Property {
				if _, ok := prop.Source().(string); ok {
					boundWidgets = append(boundWidgets, w.(Widget))
					break
				}
			}

			return true
		})

		db.SetBoundWidgets(boundWidgets)
	}
}

func (cb *ContainerBase) forEachPersistableChild(f func(p Persistable) error) error {
	if cb.children == nil {
		return nil
	}

	for _, child := range cb.children.items {
		if persistable, ok := child.(Persistable); ok && persistable.Persistent() {
			if err := f(persistable); err != nil {
				return err
			}
		}
	}

	return nil
}

func (cb *ContainerBase) Persistent() bool {
	return cb.persistent
}

func (cb *ContainerBase) SetPersistent(value bool) {
	cb.persistent = value
}

func (cb *ContainerBase) SaveState() error {
	return cb.forEachPersistableChild(func(p Persistable) error {
		return p.SaveState()
	})
}

func (cb *ContainerBase) RestoreState() error {
	return cb.forEachPersistableChild(func(p Persistable) error {
		return p.RestoreState()
	})
}

func (cb *ContainerBase) SetSuspended(suspend bool) {
	wasSuspended := cb.Suspended()

	cb.WidgetBase.SetSuspended(suspend)

	if !suspend && wasSuspended && cb.layout != nil {
		cb.layout.Update(false)
	}
}

func (cb *ContainerBase) WndProc(hwnd win.HWND, msg uint32, wParam, lParam uintptr) uintptr {
	switch msg {
	case win.WM_COMMAND:
		if lParam == 0 {
			switch win.HIWORD(uint32(wParam)) {
			case 0:
				cmdId := win.LOWORD(uint32(wParam))
				switch cmdId {
				case win.IDOK, win.IDCANCEL:
					form := ancestor(cb)
					if form == nil {
						break
					}

					dlg, ok := form.(dialogish)
					if !ok {
						break
					}

					var button *PushButton
					if cmdId == win.IDOK {
						button = dlg.DefaultButton()
					} else {
						button = dlg.CancelButton()
					}

					if button != nil && button.Visible() && button.Enabled() {
						button.raiseClicked()
					}

					break
				}

				// Menu
				actionId := uint16(win.LOWORD(uint32(wParam)))
				if action, ok := actionsById[actionId]; ok {
					action.raiseTriggered()
					return 0
				}

			case 1:
				// Accelerator
			}
		} else {
			// The window that sent the notification shall handle it itself.
			hWnd := win.HWND(lParam)
			if window := windowFromHandle(hWnd); window != nil {
				window.WndProc(hwnd, msg, wParam, lParam)
				return 0
			}
		}

	case win.WM_NOTIFY:
		nmh := (*win.NMHDR)(unsafe.Pointer(lParam))
		if window := windowFromHandle(nmh.HwndFrom); window != nil {
			// The window that sent the notification shall handle it itself.
			return window.WndProc(hwnd, msg, wParam, lParam)
		}

	case win.WM_SIZE, win.WM_SIZING:
		if cb.layout != nil {
			cb.layout.Update(false)
		}
	}

	return cb.WidgetBase.WndProc(hwnd, msg, wParam, lParam)
}

func (cb *ContainerBase) onInsertingWidget(index int, widget Widget) (err error) {
	return nil
}

func (cb *ContainerBase) onInsertedWidget(index int, widget Widget) (err error) {
	if parent := widget.Parent(); parent == nil || parent.Handle() != cb.hWnd {
		if err = widget.SetParent(cb.window.(Container)); err != nil {
			return
		}
	}

	if cb.layout != nil {
		cb.layout.Update(true)
	}

	widget.(applyFonter).applyFont(cb.Font())

	return
}

func (cb *ContainerBase) onRemovingWidget(index int, widget Widget) (err error) {
	if widget.Parent() == nil {
		return
	}

	if widget.Parent().Handle() == cb.hWnd {
		err = widget.SetParent(nil)
	}

	return
}

func (cb *ContainerBase) onRemovedWidget(index int, widget Widget) (err error) {
	if cb.layout != nil {
		cb.layout.Update(true)
	}

	return
}

func (cb *ContainerBase) onClearingWidgets() (err error) {
	for _, widget := range cb.children.items {
		if widget.Parent().Handle() == cb.hWnd {
			if err = widget.SetParent(nil); err != nil {
				return
			}
		}
	}

	return
}

func (cb *ContainerBase) onClearedWidgets() (err error) {
	if cb.layout != nil {
		cb.layout.Update(true)
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/kless/term/readline"
	"github.com/mitchellh/go-homedir"
)

var commands map[string]func(...string)

var connections []*websocket.Conn

func init() {
	commands = map[string]func(...string){
		"?":          showHelp,
		"help":       showHelp,
		"connect":    connect,
		"disconnect": disconnect,
		"send":       send,
	}
}

func main() {
	fmt.Println(`
Welcome to the websocket client. Enter ? or help for the available
commands. Press ^D (ctrl-D) to exit.
`)

	home, err := homedir.Dir()
	if err != nil {
		log.Fatalf("wsclient: failed to read home directory: %v", err)
	}
	hist, err := readline.NewHistory(filepath.Join(home, ".wsclienthist"))
	if err != nil {
		log.Fatalf("wsclient: failed to create history file: %v", err)
	}

	line, err := readline.NewDefaultLine(hist)
	if err != nil {
		log.Fatalf("wsclient: failed to create line: %v", err)
	}

	go func() {
		for {
			select {
			case <-readline.ChanCtrlC:
			case <-readline.ChanCtrlD:
				line.Restore()
				for _, conn := range connections {
					if conn != nil {
						conn.Close()
					}
				}
				os.Exit(0)
			}
		}
	}()

	for {
		l, err := line.Read()
		if err != nil {
			line.Restore()
			log.Fatalf("wsclient: failed to read line: %v", err)
		}
		args := strings.Fields(l)
		if len(args) != 0 {
			if cmd := commands[args[0]]; cmd != nil {
				cmd(args[1:]...)
			} else {
				fmt.Printf("unknown command %q\r\n", args[0])
			}
		}
	}
}

func showHelp(_ ...string) {
	keys := make([]string, 0, len(commands))
	for k := range commands {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		fmt.Printf("%s\r\n", k)
	}
}

func connect(args ...string) {
	var d websocket.Dialer

	if len(args) < 1 {
		fmt.Print("usage: connect URL [PROTO]\r\n")
		return
	}

	var h http.Header
	if len(args) == 2 {
		h = http.Header{"Sec-WebSocket-Protocol": {args[1]}}
	}
	conn, _, err := d.Dial(args[0], h)
	if err != nil {
		fmt.Printf("error: %v\r\n", err)
		return
	}
	connections = append(connections, conn)
	fmt.Printf("connected to %s [%d]\r\n", args[0], len(connections))
	go read(len(connections), conn)
}

func read(ix int, c *websocket.Conn) {
	for {
		_, b, err := c.ReadMessage()
		if err != nil {
			fmt.Printf("[%d] NextReader failed: %v; closing connection\r\n", ix, err)
			c.Close()
			return
		}
		fmt.Printf("[%d] %v\r\n", ix, string(b))
	}
}

func getConn(arg string) (*websocket.Conn, int) {
	ix, err := strconv.Atoi(arg)
	if err != nil {
		fmt.Printf("argument error: %v\r\n", err)
		return nil, 0
	}
	if ix > 0 && ix <= len(connections) {
		if c := connections[ix-1]; c != nil {
			return c, ix - 1
		}
	}
	return nil, 0
}

func disconnect(args ...string) {
	if len(args) != 1 {
		fmt.Print("usage: disconnect CONN_ID\r\n")
		return
	}
	if c, ix := getConn(args[0]); c != nil {
		c.Close()
		connections[ix] = nil
	}
}

func send(args ...string) {
	if len(args) < 2 {
		fmt.Print("usage: send CONN_ID MSG\r\n")
		return
	}
	if c, _ := getConn(args[0]); c != nil {
		if err := c.WriteMessage(websocket.TextMessage, []byte(strings.Join(args[1:], " "))); err != nil {
			fmt.Printf("WriteMessage failed: %v\r\n", err)
			return
		}
	}
}
<commit_msg>juggler/wsclient: save history on exit<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/kless/term/readline"
	"github.com/mitchellh/go-homedir"
)

var commands map[string]func(...string)

var connections []*websocket.Conn

func init() {
	commands = map[string]func(...string){
		"?":          showHelp,
		"help":       showHelp,
		"connect":    connect,
		"disconnect": disconnect,
		"send":       send,
	}
}

func main() {
	fmt.Println(`
Welcome to the websocket client. Enter ? or help for the available
commands. Press ^D (ctrl-D) to exit.
`)

	home, err := homedir.Dir()
	if err != nil {
		log.Fatalf("wsclient: failed to read home directory: %v", err)
	}
	hist, err := readline.NewHistory(filepath.Join(home, ".wsclienthist"))
	if err != nil {
		log.Fatalf("wsclient: failed to create history file: %v", err)
	}

	line, err := readline.NewDefaultLine(hist)
	if err != nil {
		log.Fatalf("wsclient: failed to create line: %v", err)
	}

	go func() {
		for {
			select {
			case <-readline.ChanCtrlC:
			case <-readline.ChanCtrlD:
				line.Restore()
				for _, conn := range connections {
					if conn != nil {
						conn.Close()
					}
				}
				hist.Save()
				os.Exit(0)
			}
		}
	}()

	for {
		l, err := line.Read()
		if err != nil {
			line.Restore()
			log.Fatalf("wsclient: failed to read line: %v", err)
		}
		args := strings.Fields(l)
		if len(args) != 0 {
			if cmd := commands[args[0]]; cmd != nil {
				cmd(args[1:]...)
			} else {
				fmt.Printf("unknown command %q\r\n", args[0])
			}
		}
	}
}

func showHelp(_ ...string) {
	keys := make([]string, 0, len(commands))
	for k := range commands {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		fmt.Printf("%s\r\n", k)
	}
}

func connect(args ...string) {
	var d websocket.Dialer

	if len(args) < 1 {
		fmt.Print("usage: connect URL [PROTO]\r\n")
		return
	}

	var h http.Header
	if len(args) == 2 {
		h = http.Header{"Sec-WebSocket-Protocol": {args[1]}}
	}
	conn, _, err := d.Dial(args[0], h)
	if err != nil {
		fmt.Printf("error: %v\r\n", err)
		return
	}
	connections = append(connections, conn)
	fmt.Printf("connected to %s [%d]\r\n", args[0], len(connections))
	go read(len(connections), conn)
}

func read(ix int, c *websocket.Conn) {
	for {
		_, b, err := c.ReadMessage()
		if err != nil {
			fmt.Printf("[%d] NextReader failed: %v; closing connection\r\n", ix, err)
			c.Close()
			return
		}
		fmt.Printf("[%d] %v\r\n", ix, string(b))
	}
}

func getConn(arg string) (*websocket.Conn, int) {
	ix, err := strconv.Atoi(arg)
	if err != nil {
		fmt.Printf("argument error: %v\r\n", err)
		return nil, 0
	}
	if ix > 0 && ix <= len(connections) {
		if c := connections[ix-1]; c != nil {
			return c, ix - 1
		}
	}
	return nil, 0
}

func disconnect(args ...string) {
	if len(args) != 1 {
		fmt.Print("usage: disconnect CONN_ID\r\n")
		return
	}
	if c, ix := getConn(args[0]); c != nil {
		c.Close()
		connections[ix] = nil
	}
}

func send(args ...string) {
	if len(args) < 2 {
		fmt.Print("usage: send CONN_ID MSG\r\n")
		return
	}
	if c, _ := getConn(args[0]); c != nil {
		if err := c.WriteMessage(websocket.TextMessage, []byte(strings.Join(args[1:], " "))); err != nil {
			fmt.Printf("WriteMessage failed: %v\r\n", err)
			return
		}
	}
}
<|endoftext|>


<commit_before>package r2dq

import (
	"os"
	"testing"
)

func tearUp() *Queue {
	redisAddr := "localhost:6379"
	env := os.Getenv("REDIS_ADDR")
	if env != "" {
		redisAddr = env
	}

	return NewQueue(redisAddr, 0, "test")
}

func tearDown(q *Queue) {
	q.redisConn.Del(q.waitingQueueKey())
	q.redisConn.Del(q.procQueueKey())
	q.Close()
}

func TestQueue(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	err := q.Queue("drteeth")
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	res := q.redisConn.LLen(q.waitingQueueKey())
	if res.Val() != 1 {
		t.Errorf("Expected %d but got %d", 1, res.Val())
	}

	q.Queue("floyd")
	res = q.redisConn.LLen(q.waitingQueueKey())
	if res.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, res.Val())
	}
}

func TestDequeue(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	q.Queue("drteeth")
	length := q.redisConn.LLen(q.waitingQueueKey())
	if length.Val() != 1 {
		t.Errorf("Expected %d but got %d", 1, length.Val())
	}

	q.Queue("floyd")
	length = q.redisConn.LLen(q.waitingQueueKey())
	if length.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, length.Val())
	}

	res, err := q.Dequeue()
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
		t.FailNow()
	}

	if res != "drteeth" {
		t.Errorf("Expected %s but got %s", "drteeth", res)
	}

	length = q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 1 {
		t.Errorf("Expected %d but got %d", 1, length.Val())
	}

	res, err = q.Dequeue()
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	if res != "floyd" {
		t.Errorf("Expected %s but got %s", "floyd", res)
	}

	length = q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, length.Val())
	}

	res, err = q.Dequeue()
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	if res != "" {
		t.Errorf("Expected empty queue, but got %s", res)
	}

}

func TestAck(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	q.Queue("drteeth")
	q.Queue("floyd")

	q.Dequeue()
	q.Dequeue()

	err := q.Ack("floyd")
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	length := q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 1 {
		t.Errorf("Expected %d but got %d", 1, length.Val())
	}

	err = q.Ack("animal")
	if err != nil {
		t.Errorf("Expected %s but got %s", ErrNotFound, err)
	}

	err = q.Ack("drteeth")
	if err != nil {
		t.Errorf("Expected %s but got %s", ErrNotFound, err)
	}

	length = q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 0 {
		t.Errorf("Expected %d but got %d", 0, length.Val())
	}
}

func TestNAck(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	q.Queue("drteeth")
	q.Queue("floyd")

	q.Dequeue()

	err := q.NAck("drteeth")
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	length := q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 0 {
		t.Errorf("Expected %d but got %d", 0, length.Val())
	}

	length = q.redisConn.LLen(q.waitingQueueKey())
	if length.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, length.Val())
	}
}

func TestGracefulShutdown(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	q.Queue("drteeth")
	q.Queue("floyd")

	q.Dequeue()
	q.Dequeue()

	q.gracefulShutdown()

	length := q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 0 {
		t.Errorf("Expected %d but got %d", 0, length.Val())
	}

	length = q.redisConn.LLen(q.waitingQueueKey())
	if length.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, length.Val())
	}

}
<commit_msg>Log TearDown errors<commit_after>package r2dq

import (
	"log"
	"os"
	"testing"
)

func tearUp() *Queue {
	redisAddr := "localhost:6379"
	env := os.Getenv("REDIS_ADDR")
	if env != "" {
		redisAddr = env
	}

	return NewQueue(redisAddr, 0, "test")
}

func tearDown(q *Queue) {
	res := q.redisConn.Del(q.waitingQueueKey())

	if res.Err() != nil {
		log.Printf("An error occurred in tearDown: %s", res.Err())
	}

	res = q.redisConn.Del(q.procQueueKey())
	if res.Err() != nil {
		log.Printf("An error occurred in tearDown: %s", res.Err())
	}

	q.Close()
}

func TestQueue(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	err := q.Queue("drteeth")
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	res := q.redisConn.LLen(q.waitingQueueKey())
	if res.Val() != 1 {
		t.Errorf("Expected %d but got %d", 1, res.Val())
	}

	q.Queue("floyd")
	res = q.redisConn.LLen(q.waitingQueueKey())
	if res.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, res.Val())
	}
}

func TestDequeue(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	q.Queue("drteeth")
	length := q.redisConn.LLen(q.waitingQueueKey())
	if length.Val() != 1 {
		t.Errorf("Expected %d but got %d", 1, length.Val())
	}

	q.Queue("floyd")
	length = q.redisConn.LLen(q.waitingQueueKey())
	if length.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, length.Val())
	}

	res, err := q.Dequeue()
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
		t.FailNow()
	}

	if res != "drteeth" {
		t.Errorf("Expected %s but got %s", "drteeth", res)
	}

	length = q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 1 {
		t.Errorf("Expected %d but got %d", 1, length.Val())
	}

	res, err = q.Dequeue()
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	if res != "floyd" {
		t.Errorf("Expected %s but got %s", "floyd", res)
	}

	length = q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, length.Val())
	}

	res, err = q.Dequeue()
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	if res != "" {
		t.Errorf("Expected empty queue, but got %s", res)
	}

}

func TestAck(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	q.Queue("drteeth")
	q.Queue("floyd")

	q.Dequeue()
	q.Dequeue()

	err := q.Ack("floyd")
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	length := q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 1 {
		t.Errorf("Expected %d but got %d", 1, length.Val())
	}

	err = q.Ack("animal")
	if err != nil {
		t.Errorf("Expected %s but got %s", ErrNotFound, err)
	}

	err = q.Ack("drteeth")
	if err != nil {
		t.Errorf("Expected %s but got %s", ErrNotFound, err)
	}

	length = q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 0 {
		t.Errorf("Expected %d but got %d", 0, length.Val())
	}
}

func TestNAck(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	q.Queue("drteeth")
	q.Queue("floyd")

	q.Dequeue()

	err := q.NAck("drteeth")
	if err != nil {
		t.Errorf("Expected nil but got %s", err)
	}

	length := q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 0 {
		t.Errorf("Expected %d but got %d", 0, length.Val())
	}

	length = q.redisConn.LLen(q.waitingQueueKey())
	if length.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, length.Val())
	}
}

func TestGracefulShutdown(t *testing.T) {
	q := tearUp()
	defer tearDown(q)

	q.Queue("drteeth")
	q.Queue("floyd")

	q.Dequeue()
	q.Dequeue()

	q.gracefulShutdown()

	length := q.redisConn.LLen(q.procQueueKey())
	if length.Val() != 0 {
		t.Errorf("Expected %d but got %d", 0, length.Val())
	}

	length = q.redisConn.LLen(q.waitingQueueKey())
	if length.Val() != 2 {
		t.Errorf("Expected %d but got %d", 2, length.Val())
	}

}
<|endoftext|>


<commit_before>package wsutil

import (
	"bufio"
	"bytes"
	"context"
	"io"
	"log"
	"net"
	"net/http"

	"github.com/gobwas/pool/pbufio"
	"github.com/gobwas/ws"
)

type DebugDialer struct {
	Dialer     *ws.Dialer
	OnRequest  func([]byte, *http.Request)
	OnResponse func([]byte, *http.Response)
}

func (d *DebugDialer) Dial(ctx context.Context, urlstr string) (conn net.Conn, br *bufio.Reader, hs ws.Handshake, err error) {
	var dialer ws.Dialer
	if d.Dialer == nil {
		dialer = ws.Dialer{}
	} else {
		dialer = *d.Dialer
	}
	var (
		rawConn net.Conn
		reqBuf  bytes.Buffer

		req  *http.Request
		reqp []byte

		res  *http.Response
		resp []byte
	)
	userWrap := dialer.WrapConn
	dialer.WrapConn = func(conn net.Conn) net.Conn {
		if userWrap != nil {
			conn = userWrap(conn)
		}
		rawConn = conn
		br = pbufio.GetReader(rawConn, 4096)

		return rwConn{
			rawConn,
			&responseLimitedReader{r: br, res: &res, resp: &resp},
			io.MultiWriter(rawConn, &reqBuf),
		}
	}

	_, _, hs, err = dialer.Dial(ctx, urlstr)
	if br.Buffered() == 0 || err != nil {
		pbufio.PutReader(br)
		br = nil
	}
	if err != nil {
		return
	}

	reqp = reqBuf.Bytes()
	reqbr := pbufio.GetReader(&reqBuf, 4096)
	defer pbufio.PutReader(reqbr)
	if req, err = http.ReadRequest(reqbr); err != nil {
		return
	}
	if onRequest := d.OnRequest; onRequest != nil {
		onRequest(reqp, req)
	}
	if onResponse := d.OnResponse; onResponse != nil {
		onResponse(resp, res)
	}

	return rawConn, br, hs, nil
}

type responseLimitedReader struct {
	r   io.Reader
	b   io.Reader
	n   int
	err error

	res  **http.Response
	resp *[]byte
}

func (r *responseLimitedReader) Read(p []byte) (n int, err error) {
	if r.b == nil {
		buf := bytes.Buffer{}
		tee := io.TeeReader(r.r, &buf)
		*r.res, r.err = http.ReadResponse(bufio.NewReader(tee), nil)

		bts := buf.Bytes()
		end := bytes.Index(bts, []byte("\r\n\r\n"))
		end += int((*r.res).ContentLength)
		end += 4
		r.b = bytes.NewReader(bts[:end])

		*r.resp = bts[:end]

		log.Printf("bts: %s", bts[:end])
	}
	if r.err != nil {
		return 0, r.err
	}
	return r.b.Read(p)
}

type rwConn struct {
	net.Conn
	r io.Reader
	w io.Writer
}

func (rwc rwConn) Read(p []byte) (int, error) {
	return rwc.r.Read(p)
}
func (rwc rwConn) Write(p []byte) (int, error) {
	return rwc.w.Write(p)
}
<commit_msg>wsutil/dialer: cleanup logs<commit_after>package wsutil

import (
	"bufio"
	"bytes"
	"context"
	"io"
	"net"
	"net/http"

	"github.com/gobwas/pool/pbufio"
	"github.com/gobwas/ws"
)

type DebugDialer struct {
	Dialer     *ws.Dialer
	OnRequest  func([]byte, *http.Request)
	OnResponse func([]byte, *http.Response)
}

func (d *DebugDialer) Dial(ctx context.Context, urlstr string) (conn net.Conn, br *bufio.Reader, hs ws.Handshake, err error) {
	var dialer ws.Dialer
	if d.Dialer == nil {
		dialer = ws.Dialer{}
	} else {
		dialer = *d.Dialer
	}
	var (
		rawConn net.Conn
		reqBuf  bytes.Buffer

		req  *http.Request
		reqp []byte

		res  *http.Response
		resp []byte
	)
	userWrap := dialer.WrapConn
	dialer.WrapConn = func(conn net.Conn) net.Conn {
		if userWrap != nil {
			conn = userWrap(conn)
		}
		rawConn = conn
		br = pbufio.GetReader(rawConn, 4096)

		return rwConn{
			rawConn,
			&responseLimitedReader{r: br, res: &res, resp: &resp},
			io.MultiWriter(rawConn, &reqBuf),
		}
	}

	_, _, hs, err = dialer.Dial(ctx, urlstr)
	if br.Buffered() == 0 || err != nil {
		pbufio.PutReader(br)
		br = nil
	}
	if err != nil {
		return
	}

	reqp = reqBuf.Bytes()
	reqbr := pbufio.GetReader(&reqBuf, 4096)
	defer pbufio.PutReader(reqbr)
	if req, err = http.ReadRequest(reqbr); err != nil {
		return
	}
	if onRequest := d.OnRequest; onRequest != nil {
		onRequest(reqp, req)
	}
	if onResponse := d.OnResponse; onResponse != nil {
		onResponse(resp, res)
	}

	return rawConn, br, hs, nil
}

type responseLimitedReader struct {
	r   io.Reader
	b   io.Reader
	n   int
	err error

	res  **http.Response
	resp *[]byte
}

func (r *responseLimitedReader) Read(p []byte) (n int, err error) {
	if r.b == nil {
		buf := bytes.Buffer{}
		tee := io.TeeReader(r.r, &buf)
		*r.res, r.err = http.ReadResponse(bufio.NewReader(tee), nil)

		bts := buf.Bytes()
		end := bytes.Index(bts, []byte("\r\n\r\n"))
		end += int((*r.res).ContentLength)
		end += 4
		r.b = bytes.NewReader(bts[:end])

		*r.resp = bts[:end]
	}
	if r.err != nil {
		return 0, r.err
	}
	return r.b.Read(p)
}

type rwConn struct {
	net.Conn
	r io.Reader
	w io.Writer
}

func (rwc rwConn) Read(p []byte) (int, error) {
	return rwc.r.Read(p)
}
func (rwc rwConn) Write(p []byte) (int, error) {
	return rwc.w.Write(p)
}
<|endoftext|>


<commit_before>package raft

import (
	"errors"
	"sort"
)

const none = -1

type messageType int

const (
	msgHup messageType = iota
	msgBeat
	msgProp
	msgApp
	msgAppResp
	msgVote
	msgVoteResp
)

var mtmap = [...]string{
	msgHup:      "msgHup",
	msgBeat:     "msgBeat",
	msgProp:     "msgProp",
	msgApp:      "msgApp",
	msgAppResp:  "msgAppResp",
	msgVote:     "msgVote",
	msgVoteResp: "msgVoteResp",
}

func (mt messageType) String() string {
	return mtmap[int(mt)]
}

var errNoLeader = errors.New("no leader")

const (
	stateFollower stateType = iota
	stateCandidate
	stateLeader
)

type stateType int

var stmap = [...]string{
	stateFollower:  "stateFollower",
	stateCandidate: "stateCandidate",
	stateLeader:    "stateLeader",
}

var stepmap = [...]stepFunc{
	stateFollower:  stepFollower,
	stateCandidate: stepCandidate,
	stateLeader:    stepLeader,
}

func (st stateType) String() string {
	return stmap[int(st)]
}

type Message struct {
	Type     messageType
	To       int
	From     int
	Term     int
	LogTerm  int
	Index    int
	PrevTerm int
	Entries  []Entry
	Commit   int
}

type index struct {
	match, next int
}

func (in *index) update(n int) {
	in.match = n
	in.next = n + 1
}

func (in *index) decr() {
	if in.next--; in.next < 1 {
		in.next = 1
	}
}

type stateMachine struct {
	id int

	// the term we are participating in at any time
	term int

	// who we voted for in term
	vote int

	// the log
	log *log

	ins map[int]*index

	state stateType

	votes map[int]bool

	msgs []Message

	// the leader id
	lead int

	// pending reconfiguration
	pendingConf bool
}

func newStateMachine(id int, peers []int) *stateMachine {
	sm := &stateMachine{id: id, log: newLog(), ins: make(map[int]*index)}
	for p := range peers {
		sm.ins[p] = &index{}
	}
	sm.reset(0)
	return sm
}

func (sm *stateMachine) poll(id int, v bool) (granted int) {
	if _, ok := sm.votes[id]; !ok {
		sm.votes[id] = v
	}
	for _, vv := range sm.votes {
		if vv {
			granted++
		}
	}
	return granted
}

// send persists state to stable storage and then sends to its mailbox.
func (sm *stateMachine) send(m Message) {
	m.From = sm.id
	m.Term = sm.term
	sm.msgs = append(sm.msgs, m)
}

// sendAppend sends RRPC, with entries to the given peer.
func (sm *stateMachine) sendAppend(to int) {
	in := sm.ins[to]
	m := Message{}
	m.Type = msgApp
	m.To = to
	m.Index = in.next - 1
	m.LogTerm = sm.log.term(in.next - 1)
	m.Entries = sm.log.entries(in.next)
	m.Commit = sm.log.committed
	sm.send(m)
}

// bcastAppend sends RRPC, with entries to all peers that are not up-to-date according to sm.mis.
func (sm *stateMachine) bcastAppend() {
	for i := range sm.ins {
		if i == sm.id {
			continue
		}
		sm.sendAppend(i)
	}
}

func (sm *stateMachine) maybeCommit() bool {
	// TODO(bmizerany): optimize.. Currently naive
	mis := make([]int, len(sm.ins))
	for i := range mis {
		mis[i] = sm.ins[i].match
	}
	sort.Sort(sort.Reverse(sort.IntSlice(mis)))
	mci := mis[sm.q()-1]

	return sm.log.maybeCommit(mci, sm.term)
}

// nextEnts returns the appliable entries and updates the applied index
func (sm *stateMachine) nextEnts() (ents []Entry) {
	return sm.log.nextEnts()
}

func (sm *stateMachine) reset(term int) {
	sm.term = term
	sm.lead = none
	sm.vote = none
	sm.votes = make(map[int]bool)
	for i := range sm.ins {
		sm.ins[i] = &index{next: sm.log.lastIndex() + 1}
		if i == sm.id {
			sm.ins[i].match = sm.log.lastIndex()
		}
	}
}

func (sm *stateMachine) q() int {
	return len(sm.ins)/2 + 1
}

func (sm *stateMachine) becomeFollower(term, lead int) {
	sm.reset(term)
	sm.lead = lead
	sm.state = stateFollower
	sm.pendingConf = false
}

func (sm *stateMachine) becomeCandidate() {
	// TODO(xiangli) remove the panic when the raft implementation is stable
	if sm.state == stateLeader {
		panic("invalid transition [leader -> candidate]")
	}
	sm.reset(sm.term + 1)
	sm.vote = sm.id
	sm.state = stateCandidate
}

func (sm *stateMachine) becomeLeader() {
	// TODO(xiangli) remove the panic when the raft implementation is stable
	if sm.state == stateFollower {
		panic("invalid transition [follower -> leader]")
	}
	sm.reset(sm.term)
	sm.lead = sm.id
	sm.state = stateLeader

	for _, e := range sm.log.ents[sm.log.committed:] {
		if e.isConfig() {
			sm.pendingConf = true
		}
	}
}

func (sm *stateMachine) Msgs() []Message {
	msgs := sm.msgs
	sm.msgs = make([]Message, 0)

	return msgs
}

func (sm *stateMachine) Step(m Message) (ok bool) {
	if m.Type == msgHup {
		sm.becomeCandidate()
		if sm.q() == sm.poll(sm.id, true) {
			sm.becomeLeader()
			return true
		}
		for i := range sm.ins {
			if i == sm.id {
				continue
			}
			lasti := sm.log.lastIndex()
			sm.send(Message{To: i, Type: msgVote, Index: lasti, LogTerm: sm.log.term(lasti)})
		}
		return true
	}

	switch {
	case m.Term == 0:
		// local message
	case m.Term > sm.term:
		sm.becomeFollower(m.Term, m.From)
	case m.Term < sm.term:
		// ignore
		return true
	}

	return stepmap[sm.state](sm, m)
}

func (sm *stateMachine) handleAppendEntries(m Message) {
	if sm.log.maybeAppend(m.Index, m.LogTerm, m.Commit, m.Entries...) {
		sm.send(Message{To: m.From, Type: msgAppResp, Index: sm.log.lastIndex()})
	} else {
		sm.send(Message{To: m.From, Type: msgAppResp, Index: -1})
	}
}

func (sm *stateMachine) addNode(id int) {
	sm.ins[id] = &index{next: sm.log.lastIndex() + 1}
	sm.pendingConf = false
}

func (sm *stateMachine) removeNode(id int) {
	delete(sm.ins, id)
	sm.pendingConf = false
}

type stepFunc func(sm *stateMachine, m Message) bool

func stepLeader(sm *stateMachine, m Message) bool {
	switch m.Type {
	case msgBeat:
		sm.bcastAppend()
	case msgProp:
		if len(m.Entries) != 1 {
			panic("unexpected length(entries) of a msgProp")
		}
		e := m.Entries[0]
		if e.isConfig() {
			if sm.pendingConf {
				return false
			}
			sm.pendingConf = true
		}
		e.Term = sm.term

		sm.log.append(sm.log.lastIndex(), e)
		sm.ins[sm.id].update(sm.log.lastIndex())
		sm.maybeCommit()
		sm.bcastAppend()
	case msgAppResp:
		if m.Index < 0 {
			sm.ins[m.From].decr()
			sm.sendAppend(m.From)
		} else {
			sm.ins[m.From].update(m.Index)
			if sm.maybeCommit() {
				sm.bcastAppend()
			}
		}
	case msgVote:
		sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
	}
	return true
}

func stepCandidate(sm *stateMachine, m Message) bool {
	switch m.Type {
	case msgProp:
		return false
	case msgApp:
		sm.becomeFollower(sm.term, m.From)
		sm.handleAppendEntries(m)
	case msgVote:
		sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
	case msgVoteResp:
		gr := sm.poll(m.From, m.Index >= 0)
		switch sm.q() {
		case gr:
			sm.becomeLeader()
			sm.bcastAppend()
		case len(sm.votes) - gr:
			sm.becomeFollower(sm.term, none)
		}
	}
	return true
}

func stepFollower(sm *stateMachine, m Message) bool {
	switch m.Type {
	case msgProp:
		if sm.lead == none {
			return false
		}
		m.To = sm.lead
		sm.send(m)
	case msgApp:
		sm.handleAppendEntries(m)
	case msgVote:
		if (sm.vote == none || sm.vote == m.From) && sm.log.isUpToDate(m.Index, m.LogTerm) {
			sm.vote = m.From
			sm.send(Message{To: m.From, Type: msgVoteResp, Index: sm.log.lastIndex()})
		} else {
			sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
		}
	}
	return true
}
<commit_msg>raft: range over sm.ins<commit_after>package raft

import (
	"errors"
	"sort"
)

const none = -1

type messageType int

const (
	msgHup messageType = iota
	msgBeat
	msgProp
	msgApp
	msgAppResp
	msgVote
	msgVoteResp
)

var mtmap = [...]string{
	msgHup:      "msgHup",
	msgBeat:     "msgBeat",
	msgProp:     "msgProp",
	msgApp:      "msgApp",
	msgAppResp:  "msgAppResp",
	msgVote:     "msgVote",
	msgVoteResp: "msgVoteResp",
}

func (mt messageType) String() string {
	return mtmap[int(mt)]
}

var errNoLeader = errors.New("no leader")

const (
	stateFollower stateType = iota
	stateCandidate
	stateLeader
)

type stateType int

var stmap = [...]string{
	stateFollower:  "stateFollower",
	stateCandidate: "stateCandidate",
	stateLeader:    "stateLeader",
}

var stepmap = [...]stepFunc{
	stateFollower:  stepFollower,
	stateCandidate: stepCandidate,
	stateLeader:    stepLeader,
}

func (st stateType) String() string {
	return stmap[int(st)]
}

type Message struct {
	Type     messageType
	To       int
	From     int
	Term     int
	LogTerm  int
	Index    int
	PrevTerm int
	Entries  []Entry
	Commit   int
}

type index struct {
	match, next int
}

func (in *index) update(n int) {
	in.match = n
	in.next = n + 1
}

func (in *index) decr() {
	if in.next--; in.next < 1 {
		in.next = 1
	}
}

type stateMachine struct {
	id int

	// the term we are participating in at any time
	term int

	// who we voted for in term
	vote int

	// the log
	log *log

	ins map[int]*index

	state stateType

	votes map[int]bool

	msgs []Message

	// the leader id
	lead int

	// pending reconfiguration
	pendingConf bool
}

func newStateMachine(id int, peers []int) *stateMachine {
	sm := &stateMachine{id: id, log: newLog(), ins: make(map[int]*index)}
	for _, p := range peers {
		sm.ins[p] = &index{}
	}
	sm.reset(0)
	return sm
}

func (sm *stateMachine) poll(id int, v bool) (granted int) {
	if _, ok := sm.votes[id]; !ok {
		sm.votes[id] = v
	}
	for _, vv := range sm.votes {
		if vv {
			granted++
		}
	}
	return granted
}

// send persists state to stable storage and then sends to its mailbox.
func (sm *stateMachine) send(m Message) {
	m.From = sm.id
	m.Term = sm.term
	sm.msgs = append(sm.msgs, m)
}

// sendAppend sends RRPC, with entries to the given peer.
func (sm *stateMachine) sendAppend(to int) {
	in := sm.ins[to]
	m := Message{}
	m.Type = msgApp
	m.To = to
	m.Index = in.next - 1
	m.LogTerm = sm.log.term(in.next - 1)
	m.Entries = sm.log.entries(in.next)
	m.Commit = sm.log.committed
	sm.send(m)
}

// bcastAppend sends RRPC, with entries to all peers that are not up-to-date according to sm.mis.
func (sm *stateMachine) bcastAppend() {
	for i := range sm.ins {
		if i == sm.id {
			continue
		}
		sm.sendAppend(i)
	}
}

func (sm *stateMachine) maybeCommit() bool {
	// TODO(bmizerany): optimize.. Currently naive
	mis := make([]int, 0, len(sm.ins))
	for i := range sm.ins {
		mis = append(mis, sm.ins[i].match)
	}
	sort.Sort(sort.Reverse(sort.IntSlice(mis)))
	mci := mis[sm.q()-1]

	return sm.log.maybeCommit(mci, sm.term)
}

// nextEnts returns the appliable entries and updates the applied index
func (sm *stateMachine) nextEnts() (ents []Entry) {
	return sm.log.nextEnts()
}

func (sm *stateMachine) reset(term int) {
	sm.term = term
	sm.lead = none
	sm.vote = none
	sm.votes = make(map[int]bool)
	for i := range sm.ins {
		sm.ins[i] = &index{next: sm.log.lastIndex() + 1}
		if i == sm.id {
			sm.ins[i].match = sm.log.lastIndex()
		}
	}
}

func (sm *stateMachine) q() int {
	return len(sm.ins)/2 + 1
}

func (sm *stateMachine) becomeFollower(term, lead int) {
	sm.reset(term)
	sm.lead = lead
	sm.state = stateFollower
	sm.pendingConf = false
}

func (sm *stateMachine) becomeCandidate() {
	// TODO(xiangli) remove the panic when the raft implementation is stable
	if sm.state == stateLeader {
		panic("invalid transition [leader -> candidate]")
	}
	sm.reset(sm.term + 1)
	sm.vote = sm.id
	sm.state = stateCandidate
}

func (sm *stateMachine) becomeLeader() {
	// TODO(xiangli) remove the panic when the raft implementation is stable
	if sm.state == stateFollower {
		panic("invalid transition [follower -> leader]")
	}
	sm.reset(sm.term)
	sm.lead = sm.id
	sm.state = stateLeader

	for _, e := range sm.log.ents[sm.log.committed:] {
		if e.isConfig() {
			sm.pendingConf = true
		}
	}
}

func (sm *stateMachine) Msgs() []Message {
	msgs := sm.msgs
	sm.msgs = make([]Message, 0)

	return msgs
}

func (sm *stateMachine) Step(m Message) (ok bool) {
	if m.Type == msgHup {
		sm.becomeCandidate()
		if sm.q() == sm.poll(sm.id, true) {
			sm.becomeLeader()
			return true
		}
		for i := range sm.ins {
			if i == sm.id {
				continue
			}
			lasti := sm.log.lastIndex()
			sm.send(Message{To: i, Type: msgVote, Index: lasti, LogTerm: sm.log.term(lasti)})
		}
		return true
	}

	switch {
	case m.Term == 0:
		// local message
	case m.Term > sm.term:
		sm.becomeFollower(m.Term, m.From)
	case m.Term < sm.term:
		// ignore
		return true
	}

	return stepmap[sm.state](sm, m)
}

func (sm *stateMachine) handleAppendEntries(m Message) {
	if sm.log.maybeAppend(m.Index, m.LogTerm, m.Commit, m.Entries...) {
		sm.send(Message{To: m.From, Type: msgAppResp, Index: sm.log.lastIndex()})
	} else {
		sm.send(Message{To: m.From, Type: msgAppResp, Index: -1})
	}
}

func (sm *stateMachine) addNode(id int) {
	sm.ins[id] = &index{next: sm.log.lastIndex() + 1}
	sm.pendingConf = false
}

func (sm *stateMachine) removeNode(id int) {
	delete(sm.ins, id)
	sm.pendingConf = false
}

type stepFunc func(sm *stateMachine, m Message) bool

func stepLeader(sm *stateMachine, m Message) bool {
	switch m.Type {
	case msgBeat:
		sm.bcastAppend()
	case msgProp:
		if len(m.Entries) != 1 {
			panic("unexpected length(entries) of a msgProp")
		}
		e := m.Entries[0]
		if e.isConfig() {
			if sm.pendingConf {
				return false
			}
			sm.pendingConf = true
		}
		e.Term = sm.term

		sm.log.append(sm.log.lastIndex(), e)
		sm.ins[sm.id].update(sm.log.lastIndex())
		sm.maybeCommit()
		sm.bcastAppend()
	case msgAppResp:
		if m.Index < 0 {
			sm.ins[m.From].decr()
			sm.sendAppend(m.From)
		} else {
			sm.ins[m.From].update(m.Index)
			if sm.maybeCommit() {
				sm.bcastAppend()
			}
		}
	case msgVote:
		sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
	}
	return true
}

func stepCandidate(sm *stateMachine, m Message) bool {
	switch m.Type {
	case msgProp:
		return false
	case msgApp:
		sm.becomeFollower(sm.term, m.From)
		sm.handleAppendEntries(m)
	case msgVote:
		sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
	case msgVoteResp:
		gr := sm.poll(m.From, m.Index >= 0)
		switch sm.q() {
		case gr:
			sm.becomeLeader()
			sm.bcastAppend()
		case len(sm.votes) - gr:
			sm.becomeFollower(sm.term, none)
		}
	}
	return true
}

func stepFollower(sm *stateMachine, m Message) bool {
	switch m.Type {
	case msgProp:
		if sm.lead == none {
			return false
		}
		m.To = sm.lead
		sm.send(m)
	case msgApp:
		sm.handleAppendEntries(m)
	case msgVote:
		if (sm.vote == none || sm.vote == m.From) && sm.log.isUpToDate(m.Index, m.LogTerm) {
			sm.vote = m.From
			sm.send(Message{To: m.From, Type: msgVoteResp, Index: sm.log.lastIndex()})
		} else {
			sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
		}
	}
	return true
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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 proxy provides a record/replay HTTP proxy. It is designed to support
// both an in-memory API (github.com/google/go-replayers/httpreplay) and a standalone server
// (github.com/google/go-replayers/httpreplay/cmd/httpr).
package proxy

// See github.com/google/martian/cmd/proxy/main.go for the origin of much of this.

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/google/martian"
	"github.com/google/martian/fifo"
	"github.com/google/martian/httpspec"
	"github.com/google/martian/martianlog"
	"github.com/google/martian/mitm"
)

// A Proxy is an HTTP proxy that supports recording or replaying requests.
type Proxy struct {
	// The certificate that the proxy uses to participate in TLS.
	CACert *x509.Certificate

	// The URL of the proxy.
	URL *url.URL

	// Initial state of the client.
	Initial []byte

	mproxy        *martian.Proxy
	filename      string          // for log
	logger        *Logger         // for recording only
	ignoreHeaders map[string]bool // headers the user has asked to ignore
}

// ForRecording returns a Proxy configured to record.
func ForRecording(filename string, port int, cert, key string) (*Proxy, error) {
	p, err := newProxy(filename, cert, key)
	if err != nil {
		return nil, err
	}

	// Construct a group that performs the standard proxy stack of request/response
	// modifications.
	stack, _ := httpspec.NewStack("httpr") // second arg is an internal group that we don't need
	p.mproxy.SetRequestModifier(stack)
	p.mproxy.SetResponseModifier(stack)

	// Make a group for logging requests and responses.
	logGroup := fifo.NewGroup()
	skipAuth := skipLoggingByHost("accounts.google.com")
	logGroup.AddRequestModifier(skipAuth)
	logGroup.AddResponseModifier(skipAuth)
	p.logger = newLogger()
	logGroup.AddRequestModifier(p.logger)
	logGroup.AddResponseModifier(p.logger)

	stack.AddRequestModifier(logGroup)
	stack.AddResponseModifier(logGroup)

	// Ordinary debug logging.
	logger := martianlog.NewLogger()
	logger.SetDecode(true)
	stack.AddRequestModifier(logger)
	stack.AddResponseModifier(logger)

	if err := p.start(port); err != nil {
		return nil, err
	}
	return p, nil
}

type hideTransport http.Transport

func (t *hideTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	return (*http.Transport)(t).RoundTrip(req)
}

var (
	configOnce sync.Once
	cert       *x509.Certificate
	config     *mitm.Config
	configErr  error
)

func newProxy(filename, c, k string) (*Proxy, error) {
	configOnce.Do(func() {
		var x509c *x509.Certificate
		var priv interface{}
		var err error
		// Set up a man-in-the-middle configuration with a CA certificate so the proxy can
		// participate in TLS.
		if c != "" && k != "" {
			x509c, priv, err = customCert(c, k)
		} else {
			x509c, priv, err = autoGenCert()
		}
		if err != nil {
			configErr = err
			return
		}
		cert = x509c
		config, configErr = mitm.NewConfig(x509c, priv)
		if config != nil {
			config.SetValidity(100 * time.Hour)
			config.SetOrganization("github.com/google/go-replayers/httpreplay")
			config.SkipTLSVerify(false)
		}
	})
	if configErr != nil {
		return nil, configErr
	}
	mproxy := martian.NewProxy()
	mproxy.SetMITM(config)
	return &Proxy{
		mproxy:        mproxy,
		CACert:        cert,
		filename:      filename,
		ignoreHeaders: map[string]bool{},
	}, nil
}

func customCert(cert, key string) (*x509.Certificate, interface{}, error) {
	tlsc, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return nil, nil, err
	}
	priv := tlsc.PrivateKey

	x509c, parseErr := x509.ParseCertificate(tlsc.Certificate[0])
	if parseErr != nil {
		return nil, nil, parseErr
	}
	return x509c, priv, nil
}

func autoGenCert() (*x509.Certificate, interface{}, error) {
	return mitm.NewAuthority("github.com/google/go-replayers/httpreplay", "HTTPReplay Authority", 100*time.Hour)
}

func (p *Proxy) start(port int) error {
	l, err := net.Listen("tcp4", fmt.Sprintf("localhost:%d", port))
	if err != nil {
		return err
	}
	p.URL = &url.URL{Scheme: "http", Host: l.Addr().String()}
	go p.mproxy.Serve(l)
	return nil
}

// Transport returns an http.Transport for clients who want to talk to the proxy.
func (p *Proxy) Transport() *http.Transport {
	caCertPool := x509.NewCertPool()
	caCertPool.AddCert(p.CACert)
	return &http.Transport{
		TLSClientConfig: &tls.Config{RootCAs: caCertPool},
		Proxy:           func(*http.Request) (*url.URL, error) { return p.URL, nil },
	}
}

// ScrubBody will replace all matching parts of the body with CLEARED.
// regexps are parsed as regexp.Regexp.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ScrubBody(regexps []string) {
	for _, re := range regexps {
		p.logger.log.Converter.registerScrubBody(re)
	}
}

// RemoveRequestHeaders will remove request headers matching patterns from the log,
// and skip matching them. Pattern is taken literally except for *, which matches any
// sequence of characters.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) RemoveRequestHeaders(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerRemoveRequestHeaders(pat)
	}
}

// ClearHeaders will replace matching headers with CLEARED.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ClearHeaders(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerClearHeaders(pat)
	}
}

// RemoveQueryParams will remove query parameters matching patterns from the request
// URL before logging, and skip matching them. Pattern is taken literally except for
// *, which matches any sequence of characters.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) RemoveQueryParams(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerRemoveParams(pat)
	}
}

// ClearQueryParams will replace matching query params in the request URL with CLEARED.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ClearQueryParams(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerClearParams(pat)
	}
}

// IgnoreHeader will cause h to be ignored during matching on replay.
// Deprecated: use RemoveRequestHeaders instead.
func (p *Proxy) IgnoreHeader(h string) {
	p.ignoreHeaders[http.CanonicalHeaderKey(h)] = true
}

// Close closes the proxy. If the proxy is recording, it also writes the log.
func (p *Proxy) Close() error {
	p.mproxy.Close()
	if p.logger != nil {
		return p.writeLog()
	}
	return nil
}

func (p *Proxy) writeLog() error {
	lg := p.logger.Extract()
	lg.Initial = p.Initial
	bytes, err := json.MarshalIndent(lg, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(p.filename, bytes, 0600) // only accessible by owner
}

// skipLoggingByHost disables logging for traffic to a particular host.
type skipLoggingByHost string

func (s skipLoggingByHost) ModifyRequest(req *http.Request) error {
	if strings.HasPrefix(req.Host, string(s)) {
		martian.NewContext(req).SkipLogging()
	}
	return nil
}

func (s skipLoggingByHost) ModifyResponse(res *http.Response) error {
	return s.ModifyRequest(res.Request)
}
<commit_msg>typed return instead of interface{}<commit_after>// Copyright 2019 Google LLC
//
// 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 proxy provides a record/replay HTTP proxy. It is designed to support
// both an in-memory API (github.com/google/go-replayers/httpreplay) and a standalone server
// (github.com/google/go-replayers/httpreplay/cmd/httpr).
package proxy

// See github.com/google/martian/cmd/proxy/main.go for the origin of much of this.

import (
	"crypto"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/google/martian"
	"github.com/google/martian/fifo"
	"github.com/google/martian/httpspec"
	"github.com/google/martian/martianlog"
	"github.com/google/martian/mitm"
)

// A Proxy is an HTTP proxy that supports recording or replaying requests.
type Proxy struct {
	// The certificate that the proxy uses to participate in TLS.
	CACert *x509.Certificate

	// The URL of the proxy.
	URL *url.URL

	// Initial state of the client.
	Initial []byte

	mproxy        *martian.Proxy
	filename      string          // for log
	logger        *Logger         // for recording only
	ignoreHeaders map[string]bool // headers the user has asked to ignore
}

// ForRecording returns a Proxy configured to record.
func ForRecording(filename string, port int, cert, key string) (*Proxy, error) {
	p, err := newProxy(filename, cert, key)
	if err != nil {
		return nil, err
	}

	// Construct a group that performs the standard proxy stack of request/response
	// modifications.
	stack, _ := httpspec.NewStack("httpr") // second arg is an internal group that we don't need
	p.mproxy.SetRequestModifier(stack)
	p.mproxy.SetResponseModifier(stack)

	// Make a group for logging requests and responses.
	logGroup := fifo.NewGroup()
	skipAuth := skipLoggingByHost("accounts.google.com")
	logGroup.AddRequestModifier(skipAuth)
	logGroup.AddResponseModifier(skipAuth)
	p.logger = newLogger()
	logGroup.AddRequestModifier(p.logger)
	logGroup.AddResponseModifier(p.logger)

	stack.AddRequestModifier(logGroup)
	stack.AddResponseModifier(logGroup)

	// Ordinary debug logging.
	logger := martianlog.NewLogger()
	logger.SetDecode(true)
	stack.AddRequestModifier(logger)
	stack.AddResponseModifier(logger)

	if err := p.start(port); err != nil {
		return nil, err
	}
	return p, nil
}

type hideTransport http.Transport

func (t *hideTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	return (*http.Transport)(t).RoundTrip(req)
}

var (
	configOnce sync.Once
	cert       *x509.Certificate
	config     *mitm.Config
	configErr  error
)

func newProxy(filename, c, k string) (*Proxy, error) {
	configOnce.Do(func() {
		var x509c *x509.Certificate
		var priv crypto.PrivateKey
		var err error
		// Set up a man-in-the-middle configuration with a CA certificate so the proxy can
		// participate in TLS.
		if c != "" && k != "" {
			x509c, priv, err = customCert(c, k)
		} else {
			x509c, priv, err = autoGenCert()
		}
		if err != nil {
			configErr = err
			return
		}
		cert = x509c
		config, configErr = mitm.NewConfig(x509c, priv)
		if config != nil {
			config.SetValidity(100 * time.Hour)
			config.SetOrganization("github.com/google/go-replayers/httpreplay")
			config.SkipTLSVerify(false)
		}
	})
	if configErr != nil {
		return nil, configErr
	}
	mproxy := martian.NewProxy()
	mproxy.SetMITM(config)
	return &Proxy{
		mproxy:        mproxy,
		CACert:        cert,
		filename:      filename,
		ignoreHeaders: map[string]bool{},
	}, nil
}

func customCert(cert, key string) (*x509.Certificate, crypto.PrivateKey, error) {
	tlsc, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return nil, nil, err
	}
	priv := tlsc.PrivateKey

	x509c, parseErr := x509.ParseCertificate(tlsc.Certificate[0])
	if parseErr != nil {
		return nil, nil, parseErr
	}
	return x509c, priv, nil
}

func autoGenCert() (*x509.Certificate, crypto.PrivateKey, error) {
	return mitm.NewAuthority("github.com/google/go-replayers/httpreplay", "HTTPReplay Authority", 100*time.Hour)
}

func (p *Proxy) start(port int) error {
	l, err := net.Listen("tcp4", fmt.Sprintf("localhost:%d", port))
	if err != nil {
		return err
	}
	p.URL = &url.URL{Scheme: "http", Host: l.Addr().String()}
	go p.mproxy.Serve(l)
	return nil
}

// Transport returns an http.Transport for clients who want to talk to the proxy.
func (p *Proxy) Transport() *http.Transport {
	caCertPool := x509.NewCertPool()
	caCertPool.AddCert(p.CACert)
	return &http.Transport{
		TLSClientConfig: &tls.Config{RootCAs: caCertPool},
		Proxy:           func(*http.Request) (*url.URL, error) { return p.URL, nil },
	}
}

// ScrubBody will replace all matching parts of the body with CLEARED.
// regexps are parsed as regexp.Regexp.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ScrubBody(regexps []string) {
	for _, re := range regexps {
		p.logger.log.Converter.registerScrubBody(re)
	}
}

// RemoveRequestHeaders will remove request headers matching patterns from the log,
// and skip matching them. Pattern is taken literally except for *, which matches any
// sequence of characters.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) RemoveRequestHeaders(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerRemoveRequestHeaders(pat)
	}
}

// ClearHeaders will replace matching headers with CLEARED.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ClearHeaders(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerClearHeaders(pat)
	}
}

// RemoveQueryParams will remove query parameters matching patterns from the request
// URL before logging, and skip matching them. Pattern is taken literally except for
// *, which matches any sequence of characters.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) RemoveQueryParams(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerRemoveParams(pat)
	}
}

// ClearQueryParams will replace matching query params in the request URL with CLEARED.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ClearQueryParams(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerClearParams(pat)
	}
}

// IgnoreHeader will cause h to be ignored during matching on replay.
// Deprecated: use RemoveRequestHeaders instead.
func (p *Proxy) IgnoreHeader(h string) {
	p.ignoreHeaders[http.CanonicalHeaderKey(h)] = true
}

// Close closes the proxy. If the proxy is recording, it also writes the log.
func (p *Proxy) Close() error {
	p.mproxy.Close()
	if p.logger != nil {
		return p.writeLog()
	}
	return nil
}

func (p *Proxy) writeLog() error {
	lg := p.logger.Extract()
	lg.Initial = p.Initial
	bytes, err := json.MarshalIndent(lg, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(p.filename, bytes, 0600) // only accessible by owner
}

// skipLoggingByHost disables logging for traffic to a particular host.
type skipLoggingByHost string

func (s skipLoggingByHost) ModifyRequest(req *http.Request) error {
	if strings.HasPrefix(req.Host, string(s)) {
		martian.NewContext(req).SkipLogging()
	}
	return nil
}

func (s skipLoggingByHost) ModifyResponse(res *http.Response) error {
	return s.ModifyRequest(res.Request)
}
<|endoftext|>


<commit_before>// +build coprocess

//go:generate msgp
//msgp:ignore CoProcessor CoProcessMiddleware CoProcessMiddlewareConfig TykMiddleware

package main

/*
#cgo python CFLAGS: -DENABLE_PYTHON
#include <stdio.h>
#include <stdlib.h>

#include "coprocess/sds/sds.h"

#include "coprocess/api.h"

#ifdef ENABLE_PYTHON
#include "coprocess/python/dispatcher.h"
#include "coprocess/python/binding.h"
#endif

*/
import "C"

import (
	"github.com/Sirupsen/logrus"
	"github.com/gorilla/context"
	"github.com/mitchellh/mapstructure"

	"github.com/TykTechnologies/tyk/coprocess"
	"github.com/TykTechnologies/tykcommon"

	"encoding/json"
	"github.com/golang/protobuf/proto"

	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
	"unsafe"
)

// EnableCoProcess will be overridden by config.EnableCoProcess.
var EnableCoProcess = false

// GlobalDispatcher will be implemented by the current CoProcess driver.
var GlobalDispatcher coprocess.Dispatcher

// CoProcessMiddleware is the basic CP middleware struct.
type CoProcessMiddleware struct {
	*TykMiddleware
	HookType         coprocess.HookType
	HookName         string
	MiddlewareDriver tykcommon.MiddlewareDriver
}

// CreateCoProcessMiddleware initializes a new CP middleware, takes hook type (pre, post, etc.), hook name ("my_hook") and driver ("python").
func CreateCoProcessMiddleware(hookName string, hookType coprocess.HookType, mwDriver tykcommon.MiddlewareDriver, tykMwSuper *TykMiddleware) func(http.Handler) http.Handler {
	dMiddleware := &CoProcessMiddleware{
		TykMiddleware:    tykMwSuper,
		HookType:         hookType,
		HookName:         hookName,
		MiddlewareDriver: mwDriver,
	}

	return CreateMiddleware(dMiddleware, tykMwSuper)
}

func doCoprocessReload() {
	log.WithFields(logrus.Fields{
		"prefix": "coprocess",
	}).Info("Reloading middlewares")
	GlobalDispatcher.Reload()

}

// CoProcessor represents a CoProcess during the request.
type CoProcessor struct {
	HookType   coprocess.HookType
	Middleware *CoProcessMiddleware
}

// GetObjectFromRequest constructs a CoProcessObject from a given http.Request.
func (c *CoProcessor) GetObjectFromRequest(r *http.Request) *coprocess.Object {

	defer r.Body.Close()
	originalBody, _ := ioutil.ReadAll(r.Body)

	var object *coprocess.Object
	var miniRequestObject *coprocess.MiniRequestObject

	miniRequestObject = &coprocess.MiniRequestObject{
		Headers:         ProtoMap(r.Header),
		SetHeaders:      make(map[string]string, 0),
		DeleteHeaders:   make([]string, 0),
		Body:            string(originalBody),
		Url:             r.URL.Path,
		Params:          ProtoMap(r.URL.Query()),
		AddParams:       make(map[string]string),
		ExtendedParams:  ProtoMap(nil),
		DeleteParams:    make([]string, 0),
		ReturnOverrides: &coprocess.ReturnOverrides{-1, ""},
	}

	object = &coprocess.Object{
		Request:  miniRequestObject,
		HookName: c.Middleware.HookName,
	}

	// If a middleware is set, take its HookType, otherwise override it with CoProcessor.HookType
	if c.Middleware != nil && c.HookType == 0 {
		c.HookType = c.Middleware.HookType
	}

	object.HookType = c.HookType

	object.Metadata = make(map[string]string, 0)
	object.Spec = make(map[string]string, 0)

	// object.Session = SessionState{}

	// Append spec data:
	if c.Middleware != nil {
		object.Spec = map[string]string{
			"OrgID": c.Middleware.TykMiddleware.Spec.OrgID,
			"APIID": c.Middleware.TykMiddleware.Spec.APIID,
		}
	}

	// Encode the session object (if not a pre-process & not a custom key check):
	if c.HookType != coprocess.HookType_Pre && c.HookType != coprocess.HookType_CustomKeyCheck {
		var session interface{}
		session = context.Get(r, SessionData)
		if session != nil {
			sessionState := session.(SessionState)
			object.Session = ProtoSessionState(sessionState)
		}
	}

	return object
}

// ObjectPostProcess does CoProcessObject post-processing (adding/removing headers or params, etc.).
func (c *CoProcessor) ObjectPostProcess(object *coprocess.Object, r *http.Request) {
	r.ContentLength = int64(len(object.Request.Body))
	r.Body = ioutil.NopCloser(bytes.NewBufferString(object.Request.Body))

	for _, dh := range object.Request.DeleteHeaders {
		r.Header.Del(dh)
	}

	for h, v := range object.Request.SetHeaders {
		r.Header.Set(h, v)
	}

	values := r.URL.Query()
	for _, k := range object.Request.DeleteParams {
		values.Del(k)
	}

	for p, v := range object.Request.AddParams {
		values.Set(p, v)
	}

	r.URL.RawQuery = values.Encode()
}

// Dispatch prepares a CoProcessMessage, sends it to the GlobalDispatcher and gets a reply.
func (c *CoProcessor) Dispatch(object *coprocess.Object) *coprocess.Object {

	var objectMsg []byte

	if MessageType == coprocess.ProtobufMessage {
		objectMsg, _ = proto.Marshal(object)
	} else if MessageType == coprocess.JsonMessage {
		objectMsg, _ = json.Marshal(object)
	}

	if CoProcessName == "grpc" {
		object = GlobalDispatcher.DispatchObject(object)
		return object
	}

	objectMsgStr := string(objectMsg)

	var CObjectStr *C.char
	CObjectStr = C.CString(objectMsgStr)

	var objectPtr *C.struct_CoProcessMessage

	objectPtr = (*C.struct_CoProcessMessage)(C.malloc(C.size_t(unsafe.Sizeof(C.struct_CoProcessMessage{}))))
	objectPtr.p_data = unsafe.Pointer(CObjectStr)
	objectPtr.length = C.int(len(objectMsg))

	var newObjectPtr *C.struct_CoProcessMessage
	newObjectPtr = (*C.struct_CoProcessMessage)(GlobalDispatcher.Dispatch(unsafe.Pointer(objectPtr)))

	var newObjectBytes []byte
	newObjectBytes = C.GoBytes(newObjectPtr.p_data, newObjectPtr.length)

	newObject := &coprocess.Object{}

	if MessageType == coprocess.ProtobufMessage {
		proto.Unmarshal(newObjectBytes, newObject)
	} else if MessageType == coprocess.JsonMessage {
		json.Unmarshal(newObjectBytes, newObject)
	}

	C.free(unsafe.Pointer(CObjectStr))
	C.free(unsafe.Pointer(objectPtr))
	C.free(unsafe.Pointer(newObjectPtr))

	return newObject
}

// CoProcessInit creates a new CoProcessDispatcher, it will be called when Tyk starts.
func CoProcessInit() (err error) {
	if config.CoProcessOptions.EnableCoProcess {
		GlobalDispatcher, err = NewCoProcessDispatcher()
		EnableCoProcess = true
	}
	return err
}

// CoProcessMiddlewareConfig holds the middleware configuration.
type CoProcessMiddlewareConfig struct {
	ConfigData map[string]string `mapstructure:"config_data" bson:"config_data" json:"config_data"`
}

// New lets you do any initialisations for the object can be done here
func (m *CoProcessMiddleware) New() {}

// GetConfig retrieves the configuration from the API config - we user mapstructure for this for simplicity
func (m *CoProcessMiddleware) GetConfig() (interface{}, error) {
	var thisModuleConfig CoProcessMiddlewareConfig

	err := mapstructure.Decode(m.TykMiddleware.Spec.APIDefinition.RawData, &thisModuleConfig)
	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "jsvm",
		}).Error(err)
		return nil, err
	}

	return thisModuleConfig, nil
}

func (m *CoProcessMiddleware) IsEnabledForSpec() bool {
	return true
}

// ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
func (m *CoProcessMiddleware) ProcessRequest(w http.ResponseWriter, r *http.Request, configuration interface{}) (error, int) {
	log.WithFields(logrus.Fields{
		"prefix": "coprocess",
	}).Debug("CoProcess Request, HookType: ", m.HookType)

	if !EnableCoProcess {
		return nil, 200
	}

	if m.HookType == coprocess.HookType_CustomKeyCheck {
		_, found := context.GetOk(r, SkipCoProcessAuth)
		if found {
			return nil, 200
		}
	}

	// It's also possible to override the HookType:
	thisCoProcessor := CoProcessor{
		Middleware: m,
		// HookType: coprocess.PreHook,
	}

	object := thisCoProcessor.GetObjectFromRequest(r)

	returnObject := thisCoProcessor.Dispatch(object)

	thisCoProcessor.ObjectPostProcess(returnObject, r)

	authHeaderValue := returnObject.Metadata["token"]

	if returnObject.Request.ReturnOverrides.ResponseCode > 400 {

		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": GetIPFromRequest(r),
			"key":    authHeaderValue,
		}).Info("Attempted access with invalid key.")

		// Fire Authfailed Event
		AuthFailed(m.TykMiddleware, r, authHeaderValue)

		// Report in health check
		ReportHealthCheckValue(m.Spec.Health, KeyFailure, "1")

		return errors.New("Key not authorised"), int(returnObject.Request.ReturnOverrides.ResponseCode)
	}

	if m.HookType == coprocess.HookType_CustomKeyCheck {
		if returnObject.Session != nil {
			var thisSessionState = TykSessionState(returnObject.Session)
			extractMiddleware := IdExtractorMiddleware{m.TykMiddleware, true, &thisSessionState}
			return extractMiddleware.ProcessRequest(w, r, configuration)
		}
	}

	// context.GetOk(r, SessionData)

	return nil, 200
}

// CoProcessLog is a bridge for using Tyk log from CP.
//export CoProcessLog
func CoProcessLog(CMessage *C.char, CLogLevel *C.char) {
	var message, logLevel string
	message = C.GoString(CMessage)
	logLevel = C.GoString(CLogLevel)

	switch logLevel {
	case "debug":
		log.WithFields(logrus.Fields{
			"prefix": CoProcessName,
		}).Debug(message)
	case "error":
		log.WithFields(logrus.Fields{
			"prefix": CoProcessName,
		}).Error(message)
	case "warning":
		log.WithFields(logrus.Fields{
			"prefix": CoProcessName,
		}).Warning(message)
	default:
		log.WithFields(logrus.Fields{
			"prefix": CoProcessName,
		}).Info(message)
	}
}
<commit_msg>Modify CP auth middleware to use ID Extractor functions.<commit_after>// +build coprocess

//go:generate msgp
//msgp:ignore CoProcessor CoProcessMiddleware CoProcessMiddlewareConfig TykMiddleware

package main

/*
#cgo python CFLAGS: -DENABLE_PYTHON
#include <stdio.h>
#include <stdlib.h>

#include "coprocess/sds/sds.h"

#include "coprocess/api.h"

#ifdef ENABLE_PYTHON
#include "coprocess/python/dispatcher.h"
#include "coprocess/python/binding.h"
#endif

*/
import "C"

import (
	"github.com/Sirupsen/logrus"
	"github.com/gorilla/context"
	"github.com/mitchellh/mapstructure"

	"github.com/TykTechnologies/tyk/coprocess"
	"github.com/TykTechnologies/tykcommon"

	"encoding/json"
	"github.com/golang/protobuf/proto"

	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
	"unsafe"
)

// EnableCoProcess will be overridden by config.EnableCoProcess.
var EnableCoProcess = false

// GlobalDispatcher will be implemented by the current CoProcess driver.
var GlobalDispatcher coprocess.Dispatcher

// CoProcessMiddleware is the basic CP middleware struct.
type CoProcessMiddleware struct {
	*TykMiddleware
	HookType         coprocess.HookType
	HookName         string
	MiddlewareDriver tykcommon.MiddlewareDriver
}

// CreateCoProcessMiddleware initializes a new CP middleware, takes hook type (pre, post, etc.), hook name ("my_hook") and driver ("python").
func CreateCoProcessMiddleware(hookName string, hookType coprocess.HookType, mwDriver tykcommon.MiddlewareDriver, tykMwSuper *TykMiddleware) func(http.Handler) http.Handler {
	dMiddleware := &CoProcessMiddleware{
		TykMiddleware:    tykMwSuper,
		HookType:         hookType,
		HookName:         hookName,
		MiddlewareDriver: mwDriver,
	}

	return CreateMiddleware(dMiddleware, tykMwSuper)
}

func doCoprocessReload() {
	log.WithFields(logrus.Fields{
		"prefix": "coprocess",
	}).Info("Reloading middlewares")
	GlobalDispatcher.Reload()

}

// CoProcessor represents a CoProcess during the request.
type CoProcessor struct {
	HookType   coprocess.HookType
	Middleware *CoProcessMiddleware
}

// GetObjectFromRequest constructs a CoProcessObject from a given http.Request.
func (c *CoProcessor) GetObjectFromRequest(r *http.Request) *coprocess.Object {

	defer r.Body.Close()
	originalBody, _ := ioutil.ReadAll(r.Body)

	var object *coprocess.Object
	var miniRequestObject *coprocess.MiniRequestObject

	miniRequestObject = &coprocess.MiniRequestObject{
		Headers:         ProtoMap(r.Header),
		SetHeaders:      make(map[string]string, 0),
		DeleteHeaders:   make([]string, 0),
		Body:            string(originalBody),
		Url:             r.URL.Path,
		Params:          ProtoMap(r.URL.Query()),
		AddParams:       make(map[string]string),
		ExtendedParams:  ProtoMap(nil),
		DeleteParams:    make([]string, 0),
		ReturnOverrides: &coprocess.ReturnOverrides{-1, ""},
	}

	object = &coprocess.Object{
		Request:  miniRequestObject,
		HookName: c.Middleware.HookName,
	}

	// If a middleware is set, take its HookType, otherwise override it with CoProcessor.HookType
	if c.Middleware != nil && c.HookType == 0 {
		c.HookType = c.Middleware.HookType
	}

	object.HookType = c.HookType

	object.Metadata = make(map[string]string, 0)
	object.Spec = make(map[string]string, 0)

	// object.Session = SessionState{}

	// Append spec data:
	if c.Middleware != nil {
		object.Spec = map[string]string{
			"OrgID": c.Middleware.TykMiddleware.Spec.OrgID,
			"APIID": c.Middleware.TykMiddleware.Spec.APIID,
		}
	}

	// Encode the session object (if not a pre-process & not a custom key check):
	if c.HookType != coprocess.HookType_Pre && c.HookType != coprocess.HookType_CustomKeyCheck {
		var session interface{}
		session = context.Get(r, SessionData)
		if session != nil {
			sessionState := session.(SessionState)
			object.Session = ProtoSessionState(sessionState)
		}
	}

	return object
}

// ObjectPostProcess does CoProcessObject post-processing (adding/removing headers or params, etc.).
func (c *CoProcessor) ObjectPostProcess(object *coprocess.Object, r *http.Request) {
	r.ContentLength = int64(len(object.Request.Body))
	r.Body = ioutil.NopCloser(bytes.NewBufferString(object.Request.Body))

	for _, dh := range object.Request.DeleteHeaders {
		r.Header.Del(dh)
	}

	for h, v := range object.Request.SetHeaders {
		r.Header.Set(h, v)
	}

	values := r.URL.Query()
	for _, k := range object.Request.DeleteParams {
		values.Del(k)
	}

	for p, v := range object.Request.AddParams {
		values.Set(p, v)
	}

	r.URL.RawQuery = values.Encode()
}

// Dispatch prepares a CoProcessMessage, sends it to the GlobalDispatcher and gets a reply.
func (c *CoProcessor) Dispatch(object *coprocess.Object) *coprocess.Object {

	var objectMsg []byte

	if MessageType == coprocess.ProtobufMessage {
		objectMsg, _ = proto.Marshal(object)
	} else if MessageType == coprocess.JsonMessage {
		objectMsg, _ = json.Marshal(object)
	}

	if CoProcessName == "grpc" {
		object = GlobalDispatcher.DispatchObject(object)
		return object
	}

	objectMsgStr := string(objectMsg)

	var CObjectStr *C.char
	CObjectStr = C.CString(objectMsgStr)

	var objectPtr *C.struct_CoProcessMessage

	objectPtr = (*C.struct_CoProcessMessage)(C.malloc(C.size_t(unsafe.Sizeof(C.struct_CoProcessMessage{}))))
	objectPtr.p_data = unsafe.Pointer(CObjectStr)
	objectPtr.length = C.int(len(objectMsg))

	var newObjectPtr *C.struct_CoProcessMessage
	newObjectPtr = (*C.struct_CoProcessMessage)(GlobalDispatcher.Dispatch(unsafe.Pointer(objectPtr)))

	var newObjectBytes []byte
	newObjectBytes = C.GoBytes(newObjectPtr.p_data, newObjectPtr.length)

	newObject := &coprocess.Object{}

	if MessageType == coprocess.ProtobufMessage {
		proto.Unmarshal(newObjectBytes, newObject)
	} else if MessageType == coprocess.JsonMessage {
		json.Unmarshal(newObjectBytes, newObject)
	}

	C.free(unsafe.Pointer(CObjectStr))
	C.free(unsafe.Pointer(objectPtr))
	C.free(unsafe.Pointer(newObjectPtr))

	return newObject
}

// CoProcessInit creates a new CoProcessDispatcher, it will be called when Tyk starts.
func CoProcessInit() (err error) {
	if config.CoProcessOptions.EnableCoProcess {
		GlobalDispatcher, err = NewCoProcessDispatcher()
		EnableCoProcess = true
	}
	return err
}

// CoProcessMiddlewareConfig holds the middleware configuration.
type CoProcessMiddlewareConfig struct {
	ConfigData map[string]string `mapstructure:"config_data" bson:"config_data" json:"config_data"`
}

// New lets you do any initialisations for the object can be done here
func (m *CoProcessMiddleware) New() {}

// GetConfig retrieves the configuration from the API config - we user mapstructure for this for simplicity
func (m *CoProcessMiddleware) GetConfig() (interface{}, error) {
	var thisModuleConfig CoProcessMiddlewareConfig

	err := mapstructure.Decode(m.TykMiddleware.Spec.APIDefinition.RawData, &thisModuleConfig)
	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "jsvm",
		}).Error(err)
		return nil, err
	}

	return thisModuleConfig, nil
}

func (m *CoProcessMiddleware) IsEnabledForSpec() bool {
	return true
}

// ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
func (m *CoProcessMiddleware) ProcessRequest(w http.ResponseWriter, r *http.Request, configuration interface{}) (error, int) {
	log.WithFields(logrus.Fields{
		"prefix": "coprocess",
	}).Debug("CoProcess Request, HookType: ", m.HookType)

	if !EnableCoProcess {
		return nil, 200
	}

	thisExtractor := m.TykMiddleware.Spec.CustomMiddleware.IdExtractor.Extractor.(IdExtractor)
	var thisSessionState *SessionState
	var returnOverrides ReturnOverrides
	var SessionID string

	if m.HookType == coprocess.HookType_CustomKeyCheck {

		SessionID, returnOverrides = thisExtractor.ExtractAndCheck(r, thisSessionState)

		if returnOverrides.ResponseCode != 0 {
			if returnOverrides.ResponseError == "" {
				return nil, returnOverrides.ResponseCode
			} else {
				err := errors.New(returnOverrides.ResponseError)
				return err, returnOverrides.ResponseCode
			}
		}
	}

	// It's also possible to override the HookType:
	thisCoProcessor := CoProcessor{
		Middleware: m,
		// HookType: coprocess.PreHook,
	}

	object := thisCoProcessor.GetObjectFromRequest(r)

	returnObject := thisCoProcessor.Dispatch(object)

	thisCoProcessor.ObjectPostProcess(returnObject, r)

	authHeaderValue := returnObject.Metadata["token"]

	if returnObject.Request.ReturnOverrides.ResponseCode > 400 {

		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": GetIPFromRequest(r),
			"key":    authHeaderValue,
		}).Info("Attempted access with invalid key.")

		// Fire Authfailed Event
		AuthFailed(m.TykMiddleware, r, authHeaderValue)

		// Report in health check
		ReportHealthCheckValue(m.Spec.Health, KeyFailure, "1")

		return errors.New("Key not authorised"), int(returnObject.Request.ReturnOverrides.ResponseCode)
	}

	if m.HookType == coprocess.HookType_CustomKeyCheck {
		if returnObject.Session != nil {
			returnedSessionState := TykSessionState(returnObject.Session)
			thisExtractor.PostProcess(r, returnedSessionState, SessionID)
			return nil, 200
		}
	}

	return nil, 200
}

// CoProcessLog is a bridge for using Tyk log from CP.
//export CoProcessLog
func CoProcessLog(CMessage *C.char, CLogLevel *C.char) {
	var message, logLevel string
	message = C.GoString(CMessage)
	logLevel = C.GoString(CLogLevel)

	switch logLevel {
	case "debug":
		log.WithFields(logrus.Fields{
			"prefix": CoProcessName,
		}).Debug(message)
	case "error":
		log.WithFields(logrus.Fields{
			"prefix": CoProcessName,
		}).Error(message)
	case "warning":
		log.WithFields(logrus.Fields{
			"prefix": CoProcessName,
		}).Warning(message)
	default:
		log.WithFields(logrus.Fields{
			"prefix": CoProcessName,
		}).Info(message)
	}
}
<|endoftext|>