<commit_before>package comm

import (
	"bufio"
	"bytes"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"
)

// Comm is some basic TCP communication
type Comm struct {
	connection net.Conn
	writer     *bufio.Writer
	reader     *bufio.Reader
}

// New returns a new comm
func New(n net.Conn) *Comm {
	c := new(Comm)
	c.connection = n
	c.connection.SetReadDeadline(time.Now().Add(3 * time.Hour))
	c.connection.SetDeadline(time.Now().Add(3 * time.Hour))
	c.connection.SetWriteDeadline(time.Now().Add(3 * time.Hour))
	c.writer = bufio.NewWriter(n)
	// c.connection = bufio.NewReader(n)
	return c
}

// Connection returns the net.TCPConn connection
func (c *Comm) Connection() net.Conn {
	return c.connection
}

// Close closes the connection
func (c *Comm) Close() {
	c.connection.Close()
}

func (c *Comm) Write(b []byte) (int, error) {
	c.writer.Write([]byte(fmt.Sprintf("%0.6d", len(b))))
	n, err := c.writer.Write(b)
	if n != len(b) {
		err = fmt.Errorf("wanted to write %d but wrote %d", n, len(b))
	}
	if err == nil {
		c.writer.Flush()
	}
	// log.Printf("wanted to write %d but wrote %d", n, len(b))
	return n, err
}

// func (c *Comm) Flush() {
// 	c.connection.Flush()
// }

func (c *Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
	// read until we get 6 bytes
	tmp := make([]byte, 6)
	n, err := c.connection.Read(tmp)
	if err != nil {
		return
	}
	tmpCopy := make([]byte, n)
	// Copy the buffer so it doesn't get changed while read by the recipient.
	copy(tmpCopy, tmp[:n])
	bs = tmpCopy

	tmp = make([]byte, 1)
	for {
		// see if we have enough bytes
		bs = bytes.Trim(bs, "\x00")
		if len(bs) == 6 {
			break
		}
		n, err := c.connection.Read(tmp)
		if err != nil {
			return nil, 0, nil, err
		}
		tmpCopy = make([]byte, n)
		// Copy the buffer so it doesn't get changed while read by the recipient.
		copy(tmpCopy, tmp[:n])
		bs = append(bs, tmpCopy...)
	}

	numBytes, err = strconv.Atoi(strings.TrimLeft(string(bs), "0"))
	if err != nil {
		return nil, 0, nil, err
	}
	buf = []byte{}
	tmp = make([]byte, numBytes)
	for {
		n, err := c.connection.Read(tmp)
		if err != nil {
			return nil, 0, nil, err
		}
		tmpCopy = make([]byte, n)
		// Copy the buffer so it doesn't get changed while read by the recipient.
		copy(tmpCopy, tmp[:n])
		buf = append(buf, bytes.TrimRight(tmpCopy, "\x00")...)
		if len(buf) < numBytes {
			// shrink the amount we need to read
			tmp = tmp[:numBytes-len(buf)]
		} else {
			break
		}
	}
	// log.Printf("wanted %d and got %d", numBytes, len(buf))
	return
}

// Send a message
func (c *Comm) Send(message string) (err error) {
	_, err = c.Write([]byte(message))
	return
}

// Receive a message
func (c *Comm) Receive() (s string, err error) {
	b, _, _, err := c.Read()
	s = string(b)
	return
}
<commit_msg>don't use writer<commit_after>package comm

import (
	"bufio"
	"bytes"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"
)

// Comm is some basic TCP communication
type Comm struct {
	connection net.Conn
	writer     *bufio.Writer
	reader     *bufio.Reader
}

// New returns a new comm
func New(n net.Conn) *Comm {
	c := new(Comm)
	c.connection = n
	c.connection.SetReadDeadline(time.Now().Add(3 * time.Hour))
	c.connection.SetDeadline(time.Now().Add(3 * time.Hour))
	c.connection.SetWriteDeadline(time.Now().Add(3 * time.Hour))
	c.writer = bufio.NewWriter(n)
	// c.connection = bufio.NewReader(n)
	return c
}

// Connection returns the net.TCPConn connection
func (c *Comm) Connection() net.Conn {
	return c.connection
}

// Close closes the connection
func (c *Comm) Close() {
	c.connection.Close()
}

func (c *Comm) Write(b []byte) (int, error) {
	c.connection.Write([]byte(fmt.Sprintf("%0.6d", len(b))))
	n, err := c.connection.Write(b)
	if n != len(b) {
		err = fmt.Errorf("wanted to write %d but wrote %d", n, len(b))
	}
	// if err == nil {
	// 	c.writer.Flush()
	// }
	// log.Printf("wanted to write %d but wrote %d", n, len(b))
	return n, err
}

// func (c *Comm) Flush() {
// 	c.connection.Flush()
// }

func (c *Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
	// read until we get 6 bytes
	tmp := make([]byte, 6)
	n, err := c.connection.Read(tmp)
	if err != nil {
		return
	}
	tmpCopy := make([]byte, n)
	// Copy the buffer so it doesn't get changed while read by the recipient.
	copy(tmpCopy, tmp[:n])
	bs = tmpCopy

	tmp = make([]byte, 1)
	for {
		// see if we have enough bytes
		bs = bytes.Trim(bs, "\x00")
		if len(bs) == 6 {
			break
		}
		n, err := c.connection.Read(tmp)
		if err != nil {
			return nil, 0, nil, err
		}
		tmpCopy = make([]byte, n)
		// Copy the buffer so it doesn't get changed while read by the recipient.
		copy(tmpCopy, tmp[:n])
		bs = append(bs, tmpCopy...)
	}

	numBytes, err = strconv.Atoi(strings.TrimLeft(string(bs), "0"))
	if err != nil {
		return nil, 0, nil, err
	}
	buf = []byte{}
	tmp = make([]byte, numBytes)
	for {
		n, err := c.connection.Read(tmp)
		if err != nil {
			return nil, 0, nil, err
		}
		tmpCopy = make([]byte, n)
		// Copy the buffer so it doesn't get changed while read by the recipient.
		copy(tmpCopy, tmp[:n])
		buf = append(buf, bytes.TrimRight(tmpCopy, "\x00")...)
		if len(buf) < numBytes {
			// shrink the amount we need to read
			tmp = tmp[:numBytes-len(buf)]
		} else {
			break
		}
	}
	// log.Printf("wanted %d and got %d", numBytes, len(buf))
	return
}

// Send a message
func (c *Comm) Send(message string) (err error) {
	_, err = c.Write([]byte(message))
	return
}

// Receive a message
func (c *Comm) Receive() (s string, err error) {
	b, _, _, err := c.Read()
	s = string(b)
	return
}
<|endoftext|>


<commit_before>package logindex

import (
	"errors"
	"github.com/howeyc/fsnotify"
	"os"
)

var (
	ErrIndexZero       = errors.New("index error")
	ErrNoFileName      = errors.New("missing file name")
	ErrNoFile          = errors.New("missing file")
	ErrInvalidFileSize = errors.New("file size invalid")
)

type Log struct {
	FileName string
	File     *os.File
	FileSize int64
	Watcher  *fsnotify.Watcher
}

type Index struct {
}

type ModPair struct {
	Last int64
	This int64
}

func (m *ModPair) delta() (delta int64, err error) {
	if m.This <= 0 || m.Last <= 0 {
		err = ErrIndexZero
		return 0, err
	}
	return (m.This - m.Last), nil
}

func New(filename string) (log *Log, err error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}

	info, err := os.Stat(filename)
	if err != nil {
		return nil, err
	}

	filesize := info.Size()

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}

	go func() {
		for {
			select {
			case ev := <-watcher.Event:
				if ev != nil && ev.IsModify() && ev.Name == filename {
				}
			case err := <-watcher.Error:
				if err != nil {
				}
			}
		}
	}()

	log = &Log{
		File:     file,
		FileName: filename,
		FileSize: filesize,
		Watcher:  watcher,
	}

	return log, nil
}

func (log *Log) watchable() (err error) {
	if log.File == nil {
		err = ErrNoFile
	}

	if log.FileName == "" {
		err = ErrNoFileName
	}

	if log.FileSize < 0 {
		err = ErrInvalidFileSize
	}

	return
}

func (log *Log) Watch() (err error) {
	err = log.watchable()
	if err != nil {
		return
	}

	log.Watcher.Watch(log.FileName)

	return
}

/*

func readAndFlush(watchfile string, pair *ModPair) {
	if pair.Last == 0 {
		pair.Last = size
	} else {
		pair.This = size

		delta := pair.delta()

		if delta > 0 {
			data := make([]byte, (delta))

			off, err := file.Seek((-1 * delta), 2)
			if err != nil {
				log.Print("Seekerr ", err)
				return
			}

			if off != 0 {
				bytesRead, err := file.Read(data)

				if err != nil {
					log.Print(err)
					return
				}
				log.Print(bytesRead, " bytes, data: ", string(data))
			}
		} else {
			return
		}
		pair.Last = size
	}
}
*/
<commit_msg>Complete refactor into more idiomatic lib.<commit_after>package logindex

import (
	"errors"
	"fmt"
	"github.com/howeyc/fsnotify"
	"os"
)

var (
	ErrIndexZero       = errors.New("index error")
	ErrNoFileName      = errors.New("missing file name")
	ErrNoFile          = errors.New("missing file")
	ErrInvalidFileSize = errors.New("file size invalid")
)

type Log struct {
	FileName string
	File     *os.File
	FileSize int64
	Watcher  *fsnotify.Watcher
	Pair     *ModPair
}

type Index struct {
}

type ModPair struct {
	Last  int64
	This  int64
	Delta int64
}

func (m *ModPair) setDelta() (err error) {
	if m.This <= 0 || m.Last <= 0 {
		err = ErrIndexZero
		return err
	}

	m.Delta = (m.This - m.Last)

	return nil
}

func New(filename string) (log *Log, err error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}

	log = &Log{
		File:     file,
		FileName: filename,
		FileSize: 0,
		Watcher:  watcher,
		Pair:     &ModPair{0, 0, 0},
	}

	if err = log.updateFileSize(); err != nil {
		return nil, err
	}

	go func() {
		for {
			select {
			case ev := <-watcher.Event:
				if ev != nil && ev.IsModify() && ev.Name == filename {
					log.moveAndFlush()
				}
			case err := <-watcher.Error:
				if err != nil {
				}
			}
		}
	}()

	return log, nil
}

func (log *Log) moveAndFlush() {
	if ok := log.movePair(); ok {
		log.flush()
	}
}

func (log *Log) movePair() (ok bool) {
	log.updateFileSize()

	if log.Pair.Last == 0 {
		log.Pair.Last = log.FileSize
		ok = false
	} else {
		log.Pair.This = log.FileSize
		ok = true
	}

	log.Pair.setDelta()

	log.Pair.Last = log.FileSize

	return
}

func (log *Log) updateFileSize() (err error) {
	info, err := os.Stat(log.FileName)
	if err != nil {
		return err
	}

	log.FileSize = info.Size()
	return nil
}

func (log *Log) flush() {
	delta := log.Pair.Delta
	file := log.File

	if delta > 0 {
		data := make([]byte, (delta))

		off, err := file.Seek((-1 * delta), 2)
		if err != nil {
			return
		}

		if off != 0 {
			bytesRead, err := file.Read(data)

			if err != nil {
				return
			}

			if bytesRead != 0 {
				fmt.Println(string(data))
			}
		}
	} else {
		return
	}

}

func (log *Log) watchable() (err error) {
	if log.File == nil {
		err = ErrNoFile
	}

	if log.FileName == "" {
		err = ErrNoFileName
	}

	if log.FileSize < 0 {
		err = ErrInvalidFileSize
	}

	return
}

func (log *Log) Watch() (err error) {
	err = log.watchable()
	if err != nil {
		return
	}

	log.Watcher.Watch(log.FileName)

	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/boivie/lovebeat-go/dashboard"
	"github.com/boivie/lovebeat-go/httpapi"
	"github.com/boivie/lovebeat-go/internal"
	"github.com/boivie/lovebeat-go/service"
	"github.com/boivie/lovebeat-go/tcpapi"
	"github.com/boivie/lovebeat-go/udpapi"
	"github.com/gorilla/mux"
	"github.com/op/go-logging"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"
)

var log = logging.MustGetLogger("lovebeat")

const (
	VERSION                 = "0.1.0"
	MAX_UNPROCESSED_PACKETS = 1000
)

var (
	udpAddr        = flag.String("udp", ":8127", "UDP service address")
	tcpAddr        = flag.String("tcp", ":8127", "TCP service address")
	expiryInterval = flag.Int64("expiry-interval", 1, "Expiry interval (seconds)")
	debug          = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion    = flag.Bool("version", false, "print version string")
)

var (
	ServiceCmdChan = make(chan *internal.Cmd, MAX_UNPROCESSED_PACKETS)
	ViewCmdChan    = make(chan *internal.ViewCmd, MAX_UNPROCESSED_PACKETS)
	signalchan     = make(chan os.Signal, 1)
)

func now() int64 { return time.Now().Unix() }

func monitor() {
	period := time.Duration(*expiryInterval) * time.Second
	ticker := time.NewTicker(period)
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			return
		case <-ticker.C:
			var ts = now()
			for _, s := range service.GetServices() {
				if s.State == service.STATE_PAUSED || s.State == s.StateAt(ts) {
					continue
				}
				var ref = *s
				s.State = s.StateAt(ts)
				s.Save(&ref, ts)
				s.UpdateViews(ViewCmdChan)
			}
		case c := <-ViewCmdChan:
			var ts = now()
			switch c.Action {
			case service.ACTION_REFRESH_VIEW:
				log.Debug("Refresh view %s", c.View)
				var view = service.GetView(c.View)
				var ref = *view
				view.Refresh(ts)
				view.Save(&ref, ts)
			}
		case c := <-ServiceCmdChan:
			var ts = now()
			var s = service.GetService(c.Service)
			var ref = *s
			switch c.Action {
			case internal.ACTION_SET_WARN:
				s.WarningTimeout = int64(c.Value)
			case internal.ACTION_SET_ERR:
				s.ErrorTimeout = int64(c.Value)
			case internal.ACTION_BEAT:
				if c.Value > 0 {
					s.LastBeat = ts
					var diff = ts - ref.LastBeat
					s.Log("%d|beat|%d", ts, diff)
					log.Debug("Beat from %s", s.Name)
				}
			}
			s.State = s.StateAt(ts)
			s.Save(&ref, ts)
			s.UpdateViews(ViewCmdChan)
		}
	}
}

func httpServer(port int16) {
	rtr := mux.NewRouter()
	httpapi.Register(rtr, ServiceCmdChan, ViewCmdChan)
	dashboard.Register(rtr)
	http.Handle("/", rtr)
	log.Info("HTTP server running on port %d\n", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
}

func main() {
	flag.Parse()

	var format = logging.MustStringFormatter("%{level} %{message}")
	logging.SetFormatter(format)
	if *debug {
		logging.SetLevel(logging.DEBUG, "lovebeat")
	} else {
		logging.SetLevel(logging.INFO, "lovebeat")
	}
	log.Debug("Debug logs enabled")

	if *showVersion {
		fmt.Printf("lovebeats v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}
	log.Info("Lovebeat v%s started as PID %d", VERSION, os.Getpid())

	service.Startup()

	signal.Notify(signalchan, syscall.SIGTERM)

	go httpServer(8080)
	go udpapi.Listener(*udpAddr, ServiceCmdChan)
	go tcpapi.Listener(*tcpAddr, ServiceCmdChan)

	log.Info("Ready to handle incoming connections")

	monitor()
}
<commit_msg>Printing hostname<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/boivie/lovebeat-go/dashboard"
	"github.com/boivie/lovebeat-go/httpapi"
	"github.com/boivie/lovebeat-go/internal"
	"github.com/boivie/lovebeat-go/service"
	"github.com/boivie/lovebeat-go/tcpapi"
	"github.com/boivie/lovebeat-go/udpapi"
	"github.com/gorilla/mux"
	"github.com/op/go-logging"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
	"time"
)

var log = logging.MustGetLogger("lovebeat")

const (
	VERSION                 = "0.1.0"
	MAX_UNPROCESSED_PACKETS = 1000
)

var (
	udpAddr        = flag.String("udp", ":8127", "UDP service address")
	tcpAddr        = flag.String("tcp", ":8127", "TCP service address")
	expiryInterval = flag.Int64("expiry-interval", 1, "Expiry interval (seconds)")
	debug          = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion    = flag.Bool("version", false, "print version string")
)

var (
	ServiceCmdChan = make(chan *internal.Cmd, MAX_UNPROCESSED_PACKETS)
	ViewCmdChan    = make(chan *internal.ViewCmd, MAX_UNPROCESSED_PACKETS)
	signalchan     = make(chan os.Signal, 1)
)

func now() int64 { return time.Now().Unix() }

func monitor() {
	period := time.Duration(*expiryInterval) * time.Second
	ticker := time.NewTicker(period)
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			return
		case <-ticker.C:
			var ts = now()
			for _, s := range service.GetServices() {
				if s.State == service.STATE_PAUSED || s.State == s.StateAt(ts) {
					continue
				}
				var ref = *s
				s.State = s.StateAt(ts)
				s.Save(&ref, ts)
				s.UpdateViews(ViewCmdChan)
			}
		case c := <-ViewCmdChan:
			var ts = now()
			switch c.Action {
			case service.ACTION_REFRESH_VIEW:
				log.Debug("Refresh view %s", c.View)
				var view = service.GetView(c.View)
				var ref = *view
				view.Refresh(ts)
				view.Save(&ref, ts)
			}
		case c := <-ServiceCmdChan:
			var ts = now()
			var s = service.GetService(c.Service)
			var ref = *s
			switch c.Action {
			case internal.ACTION_SET_WARN:
				s.WarningTimeout = int64(c.Value)
			case internal.ACTION_SET_ERR:
				s.ErrorTimeout = int64(c.Value)
			case internal.ACTION_BEAT:
				if c.Value > 0 {
					s.LastBeat = ts
					var diff = ts - ref.LastBeat
					s.Log("%d|beat|%d", ts, diff)
					log.Debug("Beat from %s", s.Name)
				}
			}
			s.State = s.StateAt(ts)
			s.Save(&ref, ts)
			s.UpdateViews(ViewCmdChan)
		}
	}
}

func httpServer(port int16) {
	rtr := mux.NewRouter()
	httpapi.Register(rtr, ServiceCmdChan, ViewCmdChan)
	dashboard.Register(rtr)
	http.Handle("/", rtr)
	log.Info("HTTP server running on port %d\n", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
}

func getHostname() string {
	var hostname, err = os.Hostname()
	if err != nil {
		return fmt.Sprintf("unknown_%d", os.Getpid())
	}
	return strings.Split(hostname, ".")[0]
}

func main() {
	flag.Parse()

	var format = logging.MustStringFormatter("%{level} %{message}")
	logging.SetFormatter(format)
	if *debug {
		logging.SetLevel(logging.DEBUG, "lovebeat")
	} else {
		logging.SetLevel(logging.INFO, "lovebeat")
	}
	log.Debug("Debug logs enabled")

	if *showVersion {
		fmt.Printf("lovebeats v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}
	var hostname = getHostname()
	log.Info("Lovebeat v%s started as host %s, PID %d", VERSION, hostname, os.Getpid())

	service.Startup()

	signal.Notify(signalchan, syscall.SIGTERM)

	go httpServer(8080)
	go udpapi.Listener(*udpAddr, ServiceCmdChan)
	go tcpapi.Listener(*tcpAddr, ServiceCmdChan)

	log.Info("Ready to handle incoming connections")

	monitor()
}
<|endoftext|>


<commit_before>package proxy

import (
	"fmt"
	"net"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/maxencoder/mixer/db"
	"github.com/maxencoder/mixer/sqlparser"
	"github.com/siddontang/go-log/log"
	. "github.com/siddontang/go-mysql/mysql"
	"github.com/siddontang/go-mysql/server"
)

//client <-> proxy
type Conn struct {
	sync.Mutex

	c *server.Conn

	server *Server

	connectionId uint32

	status uint16

	charset string

	db string

	schema *Schema

	txConns map[*Node]*db.SqlConn

	closed bool

	lastInsertId int64
	affectedRows int64
}

var baseConnID uint32 = 10000

func (s *Server) newConn(co net.Conn) (c *Conn, err error) {
	c = new(Conn)

	c.connectionId = atomic.AddUint32(&baseConnID, 1)

	c.server = s

	//c.status = SERVER_STATUS_AUTOCOMMIT // don't see this in new code?

	c.txConns = make(map[*Node]*db.SqlConn)

	c.closed = false

	c.charset = DEFAULT_CHARSET

	c.c, err = server.NewConn(co, s.user, s.password, c)

	return
}

func (c *Conn) Close() error {
	if c.closed {
		return nil
	}

	c.c.Close()

	c.rollback()

	c.closed = true

	return nil
}

func (c *Conn) Run() {
	defer func() {
		r := recover()
		if err, ok := r.(error); ok {
			const size = 4096
			buf := make([]byte, size)
			buf = buf[:runtime.Stack(buf, false)]

			log.Error("%v, %s", err, buf)
		}

		c.Close()
	}()

	for {
		err := c.c.HandleCommand()
		if err != nil {
			log.Error("dispatch error %s", err.Error())
			return
		}

		if c.c.Closed() {
			return
		}
	}
}

//
//  server.Handler interface
//

func (c *Conn) UseDB(db string) error {
	if s := c.server.getSchema(db); s == nil {
		return NewDefaultError(ER_BAD_DB_ERROR, db)
	} else {
		c.schema = s
		c.db = db
	}
	return nil
}

func (c *Conn) HandleQuery(sql string) (*Result, error) {
	return c.handleQuery(sql)
}

func (c *Conn) HandleFieldList(table string, fieldWildcard string) ([]*Field, error) {
	if c.schema == nil {
		return nil, NewDefaultError(ER_NO_DB_ERROR)
	}

	nodeName := c.schema.rule.GetRule(table).Nodes[0]

	n := c.server.getNode(nodeName)

	co, err := n.getMasterConn()
	if err != nil {
		return nil, err
	}
	defer co.Close()

	if err = co.UseDB(c.schema.db); err != nil {
		return nil, err
	}

	var fs []*Field

	if fs, err = co.FieldList(table, fieldWildcard); err != nil {
		return nil, err
	}

	return fs, nil
}

func (c *Conn) HandleStmtPrepare(query string) (params int, columns int, context interface{}, err error) {
	return c.handleStmtPrepare(query)
}

func (c *Conn) HandleStmtExecute(context interface{}, sql string, args []interface{}) (r *Result, err error) {
	s := context.(sqlparser.Statement)

	switch stmt := s.(type) {
	case *sqlparser.Select:
		r, err = c.handleSelect(stmt, sql, args)
	case *sqlparser.Insert:
		r, err = c.handleExec(s, sql, args)
	case *sqlparser.Update:
		r, err = c.handleExec(s, sql, args)
	case *sqlparser.Delete:
		r, err = c.handleExec(s, sql, args)
	case *sqlparser.Replace:
		r, err = c.handleExec(s, sql, args)
	default:
		err = fmt.Errorf("command %T not supported now", stmt)
	}

	return
}

func (c *Conn) HandleStmtClose(context interface{}) error {
	return nil
}

func (c *Conn) handleStmtPrepare(sql string) (int, int, interface{}, error) {
	if c.schema == nil {
		return 0, 0, nil, NewDefaultError(ER_NO_DB_ERROR)
	}

	sql = strings.TrimRight(sql, ";")

	p, err := sqlparser.Parse(sql)
	if err != nil {
		return 0, 0, nil, fmt.Errorf(`parse sql "%s" error`, sql)
	}

	var tableName string
	switch t := p.(type) {
	case *sqlparser.Select:
		tableName = nstring(t.From)
	case *sqlparser.Insert:
		tableName = nstring(t.Table)
	case *sqlparser.Update:
		tableName = nstring(t.Table)
	case *sqlparser.Delete:
		tableName = nstring(t.Table)
	case *sqlparser.Replace:
		tableName = nstring(t.Table)
	default:
		return 0, 0, nil, fmt.Errorf(`unsupport prepare sql "%s"`, sql)
	}

	r := c.schema.rule.GetRule(tableName)

	n := c.server.getNode(r.Nodes[0])

	var params, columns int
	if co, err := n.getMasterConn(); err != nil {
		return 0, 0, nil, fmt.Errorf("prepare error %s", err)
	} else {
		defer co.Close()

		if err = co.UseDB(c.schema.db); err != nil {
			return 0, 0, nil, fmt.Errorf("parepre error %s", err)
		}

		if t, err := co.Prepare(sql); err != nil {
			return 0, 0, nil, fmt.Errorf("parepre error %s", err)
		} else {
			params = t.ParamNum()
			columns = t.ColumnNum()
		}
	}

	return params, columns, p, nil
}
<commit_msg>not very reliable without autocommit<commit_after>package proxy

import (
	"fmt"
	"net"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/maxencoder/mixer/db"
	"github.com/maxencoder/mixer/sqlparser"
	"github.com/siddontang/go-log/log"
	. "github.com/siddontang/go-mysql/mysql"
	"github.com/siddontang/go-mysql/server"
)

//client <-> proxy
type Conn struct {
	sync.Mutex

	c *server.Conn

	server *Server

	connectionId uint32

	status uint16

	charset string

	db string

	schema *Schema

	txConns map[*Node]*db.SqlConn

	closed bool

	lastInsertId int64
	affectedRows int64
}

var baseConnID uint32 = 10000

func (s *Server) newConn(co net.Conn) (c *Conn, err error) {
	c = new(Conn)

	c.connectionId = atomic.AddUint32(&baseConnID, 1)

	c.server = s

	c.status = SERVER_STATUS_AUTOCOMMIT

	c.txConns = make(map[*Node]*db.SqlConn)

	c.closed = false

	c.charset = DEFAULT_CHARSET

	c.c, err = server.NewConn(co, s.user, s.password, c)

	return
}

func (c *Conn) Close() error {
	if c.closed {
		return nil
	}

	c.c.Close()

	c.rollback()

	c.closed = true

	return nil
}

func (c *Conn) Run() {
	defer func() {
		r := recover()
		if err, ok := r.(error); ok {
			const size = 4096
			buf := make([]byte, size)
			buf = buf[:runtime.Stack(buf, false)]

			log.Error("%v, %s", err, buf)
		}

		c.Close()
	}()

	for {
		err := c.c.HandleCommand()
		if err != nil {
			log.Error("dispatch error %s", err.Error())
			return
		}

		if c.c.Closed() {
			return
		}
	}
}

//
//  server.Handler interface
//

func (c *Conn) UseDB(db string) error {
	if s := c.server.getSchema(db); s == nil {
		return NewDefaultError(ER_BAD_DB_ERROR, db)
	} else {
		c.schema = s
		c.db = db
	}
	return nil
}

func (c *Conn) HandleQuery(sql string) (*Result, error) {
	return c.handleQuery(sql)
}

func (c *Conn) HandleFieldList(table string, fieldWildcard string) ([]*Field, error) {
	if c.schema == nil {
		return nil, NewDefaultError(ER_NO_DB_ERROR)
	}

	nodeName := c.schema.rule.GetRule(table).Nodes[0]

	n := c.server.getNode(nodeName)

	co, err := n.getMasterConn()
	if err != nil {
		return nil, err
	}
	defer co.Close()

	if err = co.UseDB(c.schema.db); err != nil {
		return nil, err
	}

	var fs []*Field

	if fs, err = co.FieldList(table, fieldWildcard); err != nil {
		return nil, err
	}

	return fs, nil
}

func (c *Conn) HandleStmtPrepare(query string) (params int, columns int, context interface{}, err error) {
	return c.handleStmtPrepare(query)
}

func (c *Conn) HandleStmtExecute(context interface{}, sql string, args []interface{}) (r *Result, err error) {
	s := context.(sqlparser.Statement)

	switch stmt := s.(type) {
	case *sqlparser.Select:
		r, err = c.handleSelect(stmt, sql, args)
	case *sqlparser.Insert:
		r, err = c.handleExec(s, sql, args)
	case *sqlparser.Update:
		r, err = c.handleExec(s, sql, args)
	case *sqlparser.Delete:
		r, err = c.handleExec(s, sql, args)
	case *sqlparser.Replace:
		r, err = c.handleExec(s, sql, args)
	default:
		err = fmt.Errorf("command %T not supported now", stmt)
	}

	return
}

func (c *Conn) HandleStmtClose(context interface{}) error {
	return nil
}

func (c *Conn) handleStmtPrepare(sql string) (int, int, interface{}, error) {
	if c.schema == nil {
		return 0, 0, nil, NewDefaultError(ER_NO_DB_ERROR)
	}

	sql = strings.TrimRight(sql, ";")

	p, err := sqlparser.Parse(sql)
	if err != nil {
		return 0, 0, nil, fmt.Errorf(`parse sql "%s" error`, sql)
	}

	var tableName string
	switch t := p.(type) {
	case *sqlparser.Select:
		tableName = nstring(t.From)
	case *sqlparser.Insert:
		tableName = nstring(t.Table)
	case *sqlparser.Update:
		tableName = nstring(t.Table)
	case *sqlparser.Delete:
		tableName = nstring(t.Table)
	case *sqlparser.Replace:
		tableName = nstring(t.Table)
	default:
		return 0, 0, nil, fmt.Errorf(`unsupport prepare sql "%s"`, sql)
	}

	r := c.schema.rule.GetRule(tableName)

	n := c.server.getNode(r.Nodes[0])

	var params, columns int
	if co, err := n.getMasterConn(); err != nil {
		return 0, 0, nil, fmt.Errorf("prepare error %s", err)
	} else {
		defer co.Close()

		if err = co.UseDB(c.schema.db); err != nil {
			return 0, 0, nil, fmt.Errorf("parepre error %s", err)
		}

		if t, err := co.Prepare(sql); err != nil {
			return 0, 0, nil, fmt.Errorf("parepre error %s", err)
		} else {
			params = t.ParamNum()
			columns = t.ColumnNum()
		}
	}

	return params, columns, p, nil
}
<|endoftext|>


<commit_before>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/netx"
	"github.com/getlantern/preconn"
	"github.com/getlantern/proxy/filters"
)

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			safeClose(downstream)
			err = errors.New("Recovered from panic handling connection: %v", p)
		}
	}()

	err = proxy.handle(ctx, downstreamIn, downstream, nil)
	return
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection: %v", p)
		}
	}()

	conn.Close()
}

func (proxy *proxy) logInitialReadError(downstream net.Conn, err error) error {
	rem := downstream.RemoteAddr()
	r := ""
	if rem != nil {
		r = rem.String()
	}
	txt := err.Error()
	// Ignore our generated error that should have already been reported.
	if strings.HasPrefix(txt, "Client Hello has no cipher suites") {
		log.Debugf("No cipher suites in common -- old Lantern client")
		return err
	}
	// These errors should all typically be internal go errors, typically with TLS. Break them up
	// for stackdriver grouping.
	if strings.Contains(txt, "oversized") {
		return log.Errorf("Oversized record on initial read: %v from %v", err, r)
	}
	if strings.Contains(txt, "first record does not") {
		return log.Errorf("Not a TLS client connection: %v from %v", err, r)
	}
	return log.Errorf("Initial ReadRequest: %v from %v", err, r)
}

func (proxy *proxy) handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn, upstream net.Conn) error {
	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
	}()

	downstreamBuffered := bufio.NewReader(downstreamIn)
	fctx := filters.WrapContext(withAwareConn(ctx), downstream)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = remoteAddr.String()
		}
		if origURLScheme(ctx) == "" {
			fctx = fctx.
				WithValue(ctxKeyOrigURLScheme, req.URL.Scheme).
				WithValue(ctxKeyOrigURLHost, req.URL.Host).
				WithValue(ctxKeyOrigHost, req.Host)
		}
	}

	if err != nil {
		if isUnexpected(err) {
			errResp := proxy.OnError(fctx, req, true, err)
			if errResp != nil {
				proxy.writeResponse(downstream, req, errResp)
			}

			return proxy.logInitialReadError(downstream, err)
		}
		return nil
	}

	var next filters.Next
	if req.Method == http.MethodConnect {
		next = proxy.nextCONNECT(downstream)
	} else {
		var tr *http.Transport
		if upstream != nil {
			setUpstreamForAwareConn(fctx, upstream)
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					// always use the supplied upstream connection, but don't allow it to
					// be closed by the transport
					return &noCloseConn{upstream}, nil
				},
				// this transport is only used once, don't keep any idle connections,
				// however still allow the transport to close the connection after using
				// it
				MaxIdleConnsPerHost: -1,
			}
		} else {
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					conn, err := proxy.Dial(ctx, false, net, addr)
					if err == nil {
						// On first dialing conn, handle RequestAware
						setUpstreamForAwareConn(ctx, conn)
						handleRequestAware(ctx)
					}
					return conn, err
				},
				IdleConnTimeout: proxy.IdleTimeout,
				// since we have one transport per downstream connection, we don't need
				// more than this
				MaxIdleConnsPerHost: 1,
			}
		}

		defer tr.CloseIdleConnections()
		next = func(ctx filters.Context, modifiedReq *http.Request) (*http.Response, filters.Context, error) {
			modifiedReq = modifiedReq.WithContext(ctx)
			setRequestForAwareConn(ctx, modifiedReq)
			handleRequestAware(ctx)
			resp, err := tr.RoundTrip(prepareRequest(modifiedReq))
			handleResponseAware(ctx, modifiedReq, resp, err)
			if err != nil {
				err = errors.New("Unable to round-trip http request to upstream: %v", err)
			}
			return resp, ctx, err
		}
	}

	return proxy.processRequests(fctx, req.RemoteAddr, req, downstream, downstreamBuffered, next)
}

func (proxy *proxy) processRequests(ctx filters.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, next filters.Next) error {
	var readErr error
	var resp *http.Response
	var err error

	for {
		if req.URL.Scheme == "" {
			req.URL.Scheme = origURLScheme(ctx)
		}
		if req.URL.Host == "" {
			req.URL.Host = origURLHost(ctx)
		}
		if req.Host == "" {
			req.Host = origHost(ctx)
		}
		resp, ctx, err = proxy.Filter.Apply(ctx, req, next)
		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, false, err)
			if resp != nil {
				// On error, we will always close the connection
				resp.Close = true
			}
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return log.Errorf("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if err != nil {
			// We encountered an error on round-tripping, stop now
			return err
		}

		upstream := upstreamConn(ctx)
		upstreamAddr := upstreamAddr(ctx)
		isConnect := upstream != nil || upstreamAddr != ""

		buffered := downstreamBuffered.Buffered()
		if buffered > 0 {
			b, _ := downstreamBuffered.Peek(buffered)
			downstream = preconn.Wrap(downstream, b)
		}

		if isConnect {
			return proxy.proceedWithConnect(ctx, req, upstreamAddr, upstream, downstream)
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				errResp := proxy.OnError(ctx, req, true, readErr)
				if errResp != nil {
					proxy.writeResponse(downstream, req, errResp)
				}
				return log.Errorf("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		ctx = ctx.IncrementRequestNumber()
		req.RemoteAddr = remoteAddr
		req = req.WithContext(ctx)
	}
}

func handleRequestAware(ctx context.Context) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case RequestAware:
			req := requestForAwareConn(ctx)
			t.OnRequest(req)
		}
		return true
	})
}

func handleResponseAware(ctx context.Context, req *http.Request, resp *http.Response, err error) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case ResponseAware:
			t.OnResponse(req, resp, err)
		}
		return true
	})
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	if resp.ProtoMajor == 0 {
		resp.ProtoMajor = 1
		resp.ProtoMinor = 1
	}
	belowHTTP11 := !resp.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
		proxy.addIdleKeepAlive(resp.Header)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	req.Close = false

	// Request Header
	newHeader := make(http.Header)
	copyHeadersForForwarding(newHeader, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	newHeader.Set("Host", req.Host)
	req.Header = newHeader

	// Request URL
	req.URL = cloneURL(req.URL)
	// If req.URL.Scheme was blank, it's http. Otherwise, it's https and we leave
	// it alone.
	if req.URL.Scheme == "" {
		req.URL.Scheme = "http"
	}
	// We need to make sure the host is defined in the URL (not the actual URI)
	req.URL.Host = req.Host

	userAgent := req.UserAgent()
	if userAgent == "" {
		req.Header.Del("User-Agent")
	} else {
		req.Header.Set("User-Agent", userAgent)
	}

	return req
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		case "Keep-Alive":
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	if err == nil {
		return false
	}
	if err == io.EOF {
		return false
	}
	// This is okay per the HTTP spec.
	// See https://www.w3.org/Protocols/rfc2616/rfc2616-sec8.html#sec8.1.4
	if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
		return false
	}

	text := err.Error()
	return !strings.HasSuffix(text, "EOF") &&
		!strings.Contains(text, "i/o timeout") &&
		!strings.Contains(text, "Use of idled network connection") &&
		!strings.Contains(text, "use of closed network connection") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "broken pipe") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "connection reset by peer")
}

func defaultFilter(ctx filters.Context, req *http.Request, next filters.Next) (*http.Response, filters.Context, error) {
	return next(ctx, req)
}

func defaultOnError(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
	return nil
}
<commit_msg>Log errors when we close HTTP connections<commit_after>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/netx"
	"github.com/getlantern/preconn"
	"github.com/getlantern/proxy/filters"
)

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			safeClose(downstream)
			err = errors.New("Recovered from panic handling connection: %v", p)
		}
	}()

	err = proxy.handle(ctx, downstreamIn, downstream, nil)
	return
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection: %v", p)
		}
	}()

	conn.Close()
}

func (proxy *proxy) logInitialReadError(downstream net.Conn, err error) error {
	rem := downstream.RemoteAddr()
	r := ""
	if rem != nil {
		r = rem.String()
	}
	txt := err.Error()
	// Ignore our generated error that should have already been reported.
	if strings.HasPrefix(txt, "Client Hello has no cipher suites") {
		log.Debugf("No cipher suites in common -- old Lantern client")
		return err
	}
	// These errors should all typically be internal go errors, typically with TLS. Break them up
	// for stackdriver grouping.
	if strings.Contains(txt, "oversized") {
		return log.Errorf("Oversized record on initial read: %v from %v", err, r)
	}
	if strings.Contains(txt, "first record does not") {
		return log.Errorf("Not a TLS client connection: %v from %v", err, r)
	}
	return log.Errorf("Initial ReadRequest: %v from %v", err, r)
}

func (proxy *proxy) handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn, upstream net.Conn) error {
	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
	}()

	downstreamBuffered := bufio.NewReader(downstreamIn)
	fctx := filters.WrapContext(withAwareConn(ctx), downstream)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = remoteAddr.String()
		}
		if origURLScheme(ctx) == "" {
			fctx = fctx.
				WithValue(ctxKeyOrigURLScheme, req.URL.Scheme).
				WithValue(ctxKeyOrigURLHost, req.URL.Host).
				WithValue(ctxKeyOrigHost, req.Host)
		}
	}

	if err != nil {
		if isUnexpected(err) {
			errResp := proxy.OnError(fctx, req, true, err)
			if errResp != nil {
				proxy.writeResponse(downstream, req, errResp)
			}

			return proxy.logInitialReadError(downstream, err)
		}
		return nil
	}

	var next filters.Next
	if req.Method == http.MethodConnect {
		next = proxy.nextCONNECT(downstream)
	} else {
		var tr *http.Transport
		if upstream != nil {
			setUpstreamForAwareConn(fctx, upstream)
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					// always use the supplied upstream connection, but don't allow it to
					// be closed by the transport
					return &noCloseConn{upstream}, nil
				},
				// this transport is only used once, don't keep any idle connections,
				// however still allow the transport to close the connection after using
				// it
				MaxIdleConnsPerHost: -1,
			}
		} else {
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					conn, err := proxy.Dial(ctx, false, net, addr)
					if err == nil {
						// On first dialing conn, handle RequestAware
						setUpstreamForAwareConn(ctx, conn)
						handleRequestAware(ctx)
					}
					return conn, err
				},
				IdleConnTimeout: proxy.IdleTimeout,
				// since we have one transport per downstream connection, we don't need
				// more than this
				MaxIdleConnsPerHost: 1,
			}
		}

		defer tr.CloseIdleConnections()
		next = func(ctx filters.Context, modifiedReq *http.Request) (*http.Response, filters.Context, error) {
			modifiedReq = modifiedReq.WithContext(ctx)
			setRequestForAwareConn(ctx, modifiedReq)
			handleRequestAware(ctx)
			resp, err := tr.RoundTrip(prepareRequest(modifiedReq))
			handleResponseAware(ctx, modifiedReq, resp, err)
			if err != nil {
				err = errors.New("Unable to round-trip http request to upstream: %v", err)
			}
			return resp, ctx, err
		}
	}

	return proxy.processRequests(fctx, req.RemoteAddr, req, downstream, downstreamBuffered, next)
}

func (proxy *proxy) processRequests(ctx filters.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, next filters.Next) error {
	var readErr error
	var resp *http.Response
	var err error

	for {
		if req.URL.Scheme == "" {
			req.URL.Scheme = origURLScheme(ctx)
		}
		if req.URL.Host == "" {
			req.URL.Host = origURLHost(ctx)
		}
		if req.Host == "" {
			req.Host = origHost(ctx)
		}
		resp, ctx, err = proxy.Filter.Apply(ctx, req, next)
		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, false, err)
			if resp != nil {
				log.Debugf("Closing client connection on error: %v", err)
				// On error, we will always close the connection
				resp.Close = true
			}
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return log.Errorf("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if err != nil {
			// We encountered an error on round-tripping, stop now
			return err
		}

		upstream := upstreamConn(ctx)
		upstreamAddr := upstreamAddr(ctx)
		isConnect := upstream != nil || upstreamAddr != ""

		buffered := downstreamBuffered.Buffered()
		if buffered > 0 {
			b, _ := downstreamBuffered.Peek(buffered)
			downstream = preconn.Wrap(downstream, b)
		}

		if isConnect {
			return proxy.proceedWithConnect(ctx, req, upstreamAddr, upstream, downstream)
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				errResp := proxy.OnError(ctx, req, true, readErr)
				if errResp != nil {
					proxy.writeResponse(downstream, req, errResp)
				}
				return log.Errorf("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		ctx = ctx.IncrementRequestNumber()
		req.RemoteAddr = remoteAddr
		req = req.WithContext(ctx)
	}
}

func handleRequestAware(ctx context.Context) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case RequestAware:
			req := requestForAwareConn(ctx)
			t.OnRequest(req)
		}
		return true
	})
}

func handleResponseAware(ctx context.Context, req *http.Request, resp *http.Response, err error) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case ResponseAware:
			t.OnResponse(req, resp, err)
		}
		return true
	})
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	if resp.ProtoMajor == 0 {
		resp.ProtoMajor = 1
		resp.ProtoMinor = 1
	}
	belowHTTP11 := !resp.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
		proxy.addIdleKeepAlive(resp.Header)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	req.Close = false

	// Request Header
	newHeader := make(http.Header)
	copyHeadersForForwarding(newHeader, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	newHeader.Set("Host", req.Host)
	req.Header = newHeader

	// Request URL
	req.URL = cloneURL(req.URL)
	// If req.URL.Scheme was blank, it's http. Otherwise, it's https and we leave
	// it alone.
	if req.URL.Scheme == "" {
		req.URL.Scheme = "http"
	}
	// We need to make sure the host is defined in the URL (not the actual URI)
	req.URL.Host = req.Host

	userAgent := req.UserAgent()
	if userAgent == "" {
		req.Header.Del("User-Agent")
	} else {
		req.Header.Set("User-Agent", userAgent)
	}

	return req
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		case "Keep-Alive":
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	if err == nil {
		return false
	}
	if err == io.EOF {
		return false
	}
	// This is okay per the HTTP spec.
	// See https://www.w3.org/Protocols/rfc2616/rfc2616-sec8.html#sec8.1.4
	if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
		return false
	}

	text := err.Error()
	return !strings.HasSuffix(text, "EOF") &&
		!strings.Contains(text, "i/o timeout") &&
		!strings.Contains(text, "Use of idled network connection") &&
		!strings.Contains(text, "use of closed network connection") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "broken pipe") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "connection reset by peer")
}

func defaultFilter(ctx filters.Context, req *http.Request, next filters.Next) (*http.Response, filters.Context, error) {
	return next(ctx, req)
}

func defaultOnError(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
	return nil
}
<|endoftext|>


<commit_before>package solidproxy

import (
	"crypto/tls"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

var (
	testMockServer *httptest.Server
)

func init() {
	// ** MOCK Server **
	handler := MockServer()
	// testMockServer = httptest.NewServer(handler)

	// testServer
	testMockServer = httptest.NewUnstartedServer(handler)
	testMockServer.TLS = new(tls.Config)
	testMockServer.TLS.ClientAuth = tls.RequestClientCert
	testMockServer.TLS.NextProtos = []string{"http/1.1"}
	testMockServer.StartTLS()
	testMockServer.URL = strings.Replace(testMockServer.URL, "127.0.0.1", "localhost", 1)
	println(testMockServer.URL)
}

func setOrigin(w http.ResponseWriter, req *http.Request) {
	origin := req.Header.Get("Origin")
	if len(origin) == 0 {
		origin = "*"
	}
	w.Header().Set("Access-Control-Allow-Origin", origin)
}

func MockServer() http.Handler {
	// Create new handler
	handler := http.NewServeMux()
	handler.Handle("/401", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		setOrigin(w, req)
		user := req.Header.Get("On-Behalf-Of")
		if len(user) == 0 {
			w.WriteHeader(401)
			w.Write([]byte("Authentication required"))
			return
		}
		if len(req.Cookies()) > 0 {
			cc := req.Cookies()[0]
			if cc.Name != "sample" && cc.Value != "sample" {
				w.WriteHeader(403)
				w.Write([]byte("Bad cookie credentials"))
				return
			}
			w.WriteHeader(200)
			w.Write([]byte("foo"))
			return
		}

		webid, err := WebIDFromReq(req)
		if err != nil {
			w.Write([]byte("\n" + err.Error()))
			return
		}
		if len(webid) > 0 {
			// set cookie
			cookie := &http.Cookie{Name: "sample", Value: "sample", HttpOnly: false}
			http.SetCookie(w, cookie)
			w.WriteHeader(200)
			w.Header().Set("User", webid)
			return
		}

		w.WriteHeader(401)
		w.Write([]byte("Authentication required"))
		return
	}))

	handler.Handle("/200", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		setOrigin(w, req)
		w.Header().Set("User-Agent-Received", req.Header.Get("User-Agent"))
		w.WriteHeader(200)
		w.Write([]byte("foo"))
		return
	}))

	handler.Handle("/method", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		setOrigin(w, req)
		w.WriteHeader(200)
		w.Write([]byte(req.Method))
		return
	}))

	return handler
}

func TestProxyMethodPOST(t *testing.T) {
	req, err := http.NewRequest("POST", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "POST", string(body))
}

func TestProxyMethodPUT(t *testing.T) {
	req, err := http.NewRequest("PUT", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "PUT", string(body))
}

func TestProxyMethodPATCH(t *testing.T) {
	req, err := http.NewRequest("PATCH", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "PATCH", string(body))
}

func TestProxyMethodDELETE(t *testing.T) {
	req, err := http.NewRequest("DELETE", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "DELETE", string(body))
}

func TestProxyMethodOPTIONS(t *testing.T) {
	req, err := http.NewRequest("OPTIONS", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "OPTIONS", string(body))
}

func TestProxyHeaders(t *testing.T) {
	origin := "example.org"
	req, err := http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/200", nil)
	assert.NoError(t, err)
	req.Header.Set("Origin", origin)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, GetServerFullName(), resp.Header.Get("User-Agent-Received"))
	assert.Equal(t, origin, resp.Header.Get("Access-Control-Allow-Origin"))
}

func TestProxyNotAuthenticated(t *testing.T) {
	req, err := http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/200", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
}

func TestProxyAuthenticated(t *testing.T) {
	alice := "https://alice.com/profile#me"

	req, err := http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 401, resp.StatusCode)

	req, err = http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	req.Header.Set("User", alice)
	resp, err = testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)

	// retry with cookie and try to remember if we have to auth from the start
	req, err = http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	req.Header.Set("User", alice)
	resp, err = testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
}

func TestProxyBadURLParse(t *testing.T) {
	req, err := http.NewRequest("GET", testProxyServer.URL+"/proxy?uri=foo", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, resp.StatusCode)

	req, err = http.NewRequest("GET", testProxyServer.URL+"/proxy?uri=http//foo.bar", nil)
	assert.NoError(t, err)
	resp, err = testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, resp.StatusCode)
}

func TestProxyBadRequest(t *testing.T) {
	req, err := http.NewRequest("FOO", testProxyServer.URL+"/proxy?uri=foo", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, resp.StatusCode)
}

func TestProxyNoSkipVerify(t *testing.T) {
	skip := false
	conf := NewServerConfig()
	conf.InsecureSkipVerify = skip
	agent, err := NewAgentLocal(testAgentWebID)
	assert.NoError(t, err)
	proxy := NewProxy(agent, skip)

	handler := NewProxyHandler(conf, proxy)
	// testProxyServer
	server := httptest.NewServer(handler)
	server.URL = strings.Replace(server.URL, "127.0.0.1", "localhost", 1)

	req, err := http.NewRequest("GET", server.URL+"/proxy?uri="+testAgentServer.URL+"/webid", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 500, resp.StatusCode)
}

func TestProxyNoUser(t *testing.T) {
	conf := NewServerConfig()
	agent, err := NewAgentLocal(testAgentWebID)
	assert.NoError(t, err)
	proxy := NewProxy(agent, true)

	handler := NewProxyHandler(conf, proxy)
	// testProxyServer
	server := httptest.NewServer(handler)
	server.URL = strings.Replace(server.URL, "127.0.0.1", "localhost", 1)

	req, err := http.NewRequest("GET", server.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 401, resp.StatusCode)
}

func TestProxyNoAgent(t *testing.T) {
	conf := NewServerConfig()
	agent, err := NewAgent(testAgentWebID)
	assert.NoError(t, err)
	proxy := NewProxy(agent, true)

	assert.Nil(t, proxy.HttpAgentClient)

	handler := NewProxyHandler(conf, proxy)
	// testProxyServer
	server := httptest.NewServer(handler)
	server.URL = strings.Replace(server.URL, "127.0.0.1", "localhost", 1)

	req, err := http.NewRequest("GET", server.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 401, resp.StatusCode)
}
<commit_msg>reordered tests<commit_after>package solidproxy

import (
	"crypto/tls"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

var (
	testMockServer *httptest.Server
)

func init() {
	// ** MOCK Server **
	handler := MockServer()
	// testMockServer = httptest.NewServer(handler)

	// testServer
	testMockServer = httptest.NewUnstartedServer(handler)
	testMockServer.TLS = new(tls.Config)
	testMockServer.TLS.ClientAuth = tls.RequestClientCert
	testMockServer.TLS.NextProtos = []string{"http/1.1"}
	testMockServer.StartTLS()
	testMockServer.URL = strings.Replace(testMockServer.URL, "127.0.0.1", "localhost", 1)
	println(testMockServer.URL)
}

func setOrigin(w http.ResponseWriter, req *http.Request) {
	origin := req.Header.Get("Origin")
	if len(origin) == 0 {
		origin = "*"
	}
	w.Header().Set("Access-Control-Allow-Origin", origin)
}

func MockServer() http.Handler {
	// Create new handler
	handler := http.NewServeMux()
	handler.Handle("/401", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		setOrigin(w, req)
		user := req.Header.Get("On-Behalf-Of")
		if len(user) == 0 {
			w.WriteHeader(401)
			w.Write([]byte("Authentication required"))
			return
		}
		if len(req.Cookies()) > 0 {
			cc := req.Cookies()[0]
			if cc.Name != "sample" && cc.Value != "sample" {
				w.WriteHeader(403)
				w.Write([]byte("Bad cookie credentials"))
				return
			}
			w.WriteHeader(200)
			w.Write([]byte("foo"))
			return
		}

		webid, err := WebIDFromReq(req)
		if err != nil {
			w.Write([]byte("\n" + err.Error()))
			return
		}
		if len(webid) > 0 {
			// set cookie
			cookie := &http.Cookie{Name: "sample", Value: "sample", HttpOnly: false}
			http.SetCookie(w, cookie)
			w.WriteHeader(200)
			w.Header().Set("User", webid)
			return
		}

		w.WriteHeader(401)
		w.Write([]byte("Authentication required"))
		return
	}))

	handler.Handle("/200", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		setOrigin(w, req)
		w.Header().Set("User-Agent-Received", req.Header.Get("User-Agent"))
		w.WriteHeader(200)
		w.Write([]byte("foo"))
		return
	}))

	handler.Handle("/method", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		setOrigin(w, req)
		w.WriteHeader(200)
		w.Write([]byte(req.Method))
		return
	}))

	return handler
}

func TestProxyMethodPOST(t *testing.T) {
	req, err := http.NewRequest("POST", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "POST", string(body))
}

func TestProxyMethodPUT(t *testing.T) {
	req, err := http.NewRequest("PUT", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "PUT", string(body))
}

func TestProxyMethodPATCH(t *testing.T) {
	req, err := http.NewRequest("PATCH", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "PATCH", string(body))
}

func TestProxyMethodDELETE(t *testing.T) {
	req, err := http.NewRequest("DELETE", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "DELETE", string(body))
}

func TestProxyMethodOPTIONS(t *testing.T) {
	req, err := http.NewRequest("OPTIONS", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/method", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	assert.Equal(t, "OPTIONS", string(body))
}

func TestProxyHeaders(t *testing.T) {
	origin := "example.org"
	req, err := http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/200", nil)
	assert.NoError(t, err)
	req.Header.Set("Origin", origin)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, GetServerFullName(), resp.Header.Get("User-Agent-Received"))
	assert.Equal(t, origin, resp.Header.Get("Access-Control-Allow-Origin"))
}

func TestProxyAuthenticated(t *testing.T) {
	alice := "https://alice.com/profile#me"

	req, err := http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	req.Header.Set("User", alice)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)

	// retry with cookie
	req, err = http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	req.Header.Set("User", alice)
	resp, err = testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
}

func TestProxyNotAuthenticated(t *testing.T) {
	req, err := http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/200", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)

	req, err = http.NewRequest("GET", testProxyServer.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	resp, err = testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 401, resp.StatusCode)
}

func TestProxyBadURLParse(t *testing.T) {
	req, err := http.NewRequest("GET", testProxyServer.URL+"/proxy?uri=foo", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, resp.StatusCode)

	req, err = http.NewRequest("GET", testProxyServer.URL+"/proxy?uri=http//foo.bar", nil)
	assert.NoError(t, err)
	resp, err = testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, resp.StatusCode)
}

func TestProxyBadRequest(t *testing.T) {
	req, err := http.NewRequest("FOO", testProxyServer.URL+"/proxy?uri=foo", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, resp.StatusCode)
}

func TestProxyNoSkipVerify(t *testing.T) {
	skip := false
	conf := NewServerConfig()
	conf.InsecureSkipVerify = skip
	agent, err := NewAgentLocal(testAgentWebID)
	assert.NoError(t, err)
	proxy := NewProxy(agent, skip)

	handler := NewProxyHandler(conf, proxy)
	// testProxyServer
	server := httptest.NewServer(handler)
	server.URL = strings.Replace(server.URL, "127.0.0.1", "localhost", 1)

	req, err := http.NewRequest("GET", server.URL+"/proxy?uri="+testAgentServer.URL+"/webid", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 500, resp.StatusCode)
}

func TestProxyNoUser(t *testing.T) {
	conf := NewServerConfig()
	agent, err := NewAgentLocal(testAgentWebID)
	assert.NoError(t, err)
	proxy := NewProxy(agent, true)

	handler := NewProxyHandler(conf, proxy)
	// testProxyServer
	server := httptest.NewServer(handler)
	server.URL = strings.Replace(server.URL, "127.0.0.1", "localhost", 1)

	req, err := http.NewRequest("GET", server.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 401, resp.StatusCode)
}

func TestProxyNoAgent(t *testing.T) {
	conf := NewServerConfig()
	agent, err := NewAgent(testAgentWebID)
	assert.NoError(t, err)
	proxy := NewProxy(agent, true)

	assert.Nil(t, proxy.HttpAgentClient)

	handler := NewProxyHandler(conf, proxy)
	// testProxyServer
	server := httptest.NewServer(handler)
	server.URL = strings.Replace(server.URL, "127.0.0.1", "localhost", 1)

	req, err := http.NewRequest("GET", server.URL+"/proxy?uri="+testMockServer.URL+"/401", nil)
	assert.NoError(t, err)
	resp, err := testClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 401, resp.StatusCode)
}
<|endoftext|>


<commit_before>// This is a library to manage Microchip MCP23008 This chip is used on onion.io Omega2 relay expension
package mcp23008

import (
	"golang.org/x/exp/io/i2c"
	"math"
	"log"
)

const (
	iodir   = 0x00
	ipol    = 0x01
	gpinten = 0x02
	defval  = 0x03
	intcon  = 0x04
	iocon   = 0x05
	gppu    = 0x06
	intf    = 0x07
	intcap  = 0x08
	gpio    = 0x09
	olat    = 0x0A
)

// McpInit function initialize MCP28003 after boot or restart of device
func McpInit(d *i2c.Device) error {
	// SetAllDirection
	err := d.WriteReg(iodir, []byte{0})
	if err != nil {
		return err
	}

	// SetAllPullUp
	err = d.WriteReg(gppu, []byte{0})
	return err
}

// McpGpioToggle change state of selected GPIO other one are unchanged
func McpGpioToggle(d *i2c.Device, gpio byte) {
	regValue := []byte{0}

	// Set 1 to corresponding BIT of GPIO
	mask := byte(math.Pow(2, float64(gpio)))

	// Read current state of all GPIO's
	d.ReadReg(gpio, regValue)

	// Write ON to selected GPIO other one keep unchanged
	d.WriteReg(gpio,[]byte{regValue[0] ^ mask})
}

// McpGpioOn set GPIO to ON/High state other one are unchanged
func McpGpioOn(d *i2c.Device, gpio byte) {
	regValue := []byte{0}

	// Set 1 to corresponding BIT of GPIO
	mask := byte(math.Pow(2, float64(gpio)))

	// Read current state of all GPIO's
	d.ReadReg(gpio, regValue)

	// Write ON to selected GPIO other one keep unchanged
	d.WriteReg(gpio,[]byte{mask | regValue[0]})
}

// Set all GPIO to ON/High state
func McpGpioAllOn(d *i2c.Device) {
	// Write ON to all GPIO
	d.WriteReg(gpio,[]byte{0xf})
}

// McpGpioOff set GPIO to OFF/Low state other one are unchanged
func McpGpioOff(d *i2c.Device, gpio byte) {
	regValue := []byte{0}

	// Set 0 to corresponding BIT of GPIO
	mask := byte(math.Pow(2, float64(gpio))) ^ 0xf

	// Read current state of all GPIO's
	d.ReadReg(gpio, regValue)

	// Write OFF to selected GPIO other one keep unchanged
	d.WriteReg(gpio,[]byte{mask & regValue[0]})
}

// Set all GPIO to OFF/Low state
func McpGpioAllOff(d *i2c.Device) {
	// Write ON to all GPIO
	d.WriteReg(gpio,[]byte{0x0})
}

// This function return state of selected GPIO 1 for ON/High or 0 for OFF/Low state
func McpReadGpio(d *i2c.Device, gpio byte) byte {
	regValue := []byte{0}

	// Set 1 to corresponding BIT of GPIO
	mask := byte(math.Pow(2, float64(gpio)))

	d.ReadReg(gpio, regValue)
	log.Printf("McpReadGpio gpio <%8b> mask <%8b> value <%8b>", gpio, mask, regValue[0])
	return (regValue[0] & mask) >> gpio
}
<commit_msg>Add log TODO: - Error Handling - Extand function - Test file<commit_after>// This is a library to manage Microchip MCP23008 This chip is used on onion.io Omega2 relay expension
package mcp23008

import (
	"golang.org/x/exp/io/i2c"
	"math"
	"log"
)

const (
	iodir   = 0x00
	ipol    = 0x01
	gpinten = 0x02
	defval  = 0x03
	intcon  = 0x04
	iocon   = 0x05
	gppu    = 0x06
	intf    = 0x07
	intcap  = 0x08
	gpio    = 0x09
	olat    = 0x0A
)

// McpInit function initialize MCP28003 after boot or restart of device
func McpInit(d *i2c.Device) error {
	// SetAllDirection
	err := d.WriteReg(iodir, []byte{0})
	if err != nil {
		return err
	}

	// SetAllPullUp
	err = d.WriteReg(gppu, []byte{0})
	return err
}

// McpGpioToggle change state of selected GPIO other one are unchanged
func McpGpioToggle(d *i2c.Device, gpio byte) {
	regValue := []byte{0}

	// Set 1 to corresponding BIT of GPIO
	mask := byte(math.Pow(2, float64(gpio)))

	// Read current state of all GPIO's
	d.ReadReg(gpio, regValue)

	// Write ON to selected GPIO other one keep unchanged
	d.WriteReg(gpio,[]byte{regValue[0] ^ mask})
}

// McpGpioOn set GPIO to ON/High state other one are unchanged
func McpGpioOn(d *i2c.Device, gpio byte) {
	regValue := []byte{0}

	// Set 1 to corresponding BIT of GPIO
	mask := byte(math.Pow(2, float64(gpio)))

	// Read current state of all GPIO's
	d.ReadReg(gpio, regValue)

	// Write ON to selected GPIO other one keep unchanged
	d.WriteReg(gpio,[]byte{mask | regValue[0]})
}

// Set all GPIO to ON/High state
func McpGpioAllOn(d *i2c.Device) {
	// Write ON to all GPIO
	d.WriteReg(gpio,[]byte{0xf})
}

// McpGpioOff set GPIO to OFF/Low state other one are unchanged
func McpGpioOff(d *i2c.Device, gpio byte) {
	regValue := []byte{0}

	// Set 0 to corresponding BIT of GPIO
	mask := byte(math.Pow(2, float64(gpio))) ^ 0xf

	// Read current state of all GPIO's
	d.ReadReg(gpio, regValue)

	// Write OFF to selected GPIO other one keep unchanged
	d.WriteReg(gpio,[]byte{mask & regValue[0]})
}

// Set all GPIO to OFF/Low state
func McpGpioAllOff(d *i2c.Device) {
	// Write ON to all GPIO
	d.WriteReg(gpio,[]byte{0x0})
}

// This function return state of selected GPIO 1 for ON/High or 0 for OFF/Low state
func McpReadGpio(d *i2c.Device, gpio byte) byte {
	regValue := []byte{0}

	// Set 1 to corresponding BIT of GPIO
	mask := byte(math.Pow(2, float64(gpio)))

	d.ReadReg(gpio, regValue)
	log.Printf("McpReadGpio gpio <%08b> mask <%08b> value <%08b>", gpio, mask, regValue[0])
	return (regValue[0] & mask) >> gpio
}
<|endoftext|>


<commit_before>package appId

import (
	"sync"

	"github.com/hashicorp/vault/helper/salt"
	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/logical/framework"
)

func Factory(conf *logical.BackendConfig) (logical.Backend, error) {
	b, err := Backend(conf)
	if err != nil {
		return nil, err
	}
	return b.Setup(conf)
}

func Backend(conf *logical.BackendConfig) (*framework.Backend, error) {
	var b backend
	b.MapAppId = &framework.PolicyMap{
		PathMap: framework.PathMap{
			Name: "app-id",
			Schema: map[string]*framework.FieldSchema{
				"display_name": &framework.FieldSchema{
					Type:        framework.TypeString,
					Description: "A name to map to this app ID for logs.",
				},

				"value": &framework.FieldSchema{
					Type:        framework.TypeString,
					Description: "Policies for the app ID.",
				},
			},
		},
		DefaultKey: "default",
	}

	b.MapUserId = &framework.PathMap{
		Name: "user-id",
		Schema: map[string]*framework.FieldSchema{
			"cidr_block": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: "If not blank, restricts auth by this CIDR block",
			},

			"value": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: "App IDs that this user associates with.",
			},
		},
	}

	b.Backend = &framework.Backend{
		Help: backendHelp,

		PathsSpecial: &logical.Paths{
			Unauthenticated: []string{
				"login",
				"login/*",
			},
		},

		Paths: framework.PathAppend([]*framework.Path{
			pathLogin(&b),
			pathLoginWithAppIDPath(&b),
		},
			b.MapAppId.Paths(),
			b.MapUserId.Paths(),
		),

		AuthRenew: b.pathLoginRenew,

		Invalidate: b.invalidate,
	}

	b.view = conf.StorageView

	return b.Backend, nil
}

type backend struct {
	*framework.Backend

	salt      *salt.Salt
	SaltMutex sync.RWMutex
	view      logical.Storage
	MapAppId  *framework.PolicyMap
	MapUserId *framework.PathMap
}

func (b *backend) Salt() (*salt.Salt, error) {
	b.SaltMutex.RLock()
	if b.salt != nil {
		defer b.SaltMutex.RUnlock()
		return b.salt, nil
	}
	b.SaltMutex.RUnlock()
	b.SaltMutex.Lock()
	defer b.SaltMutex.Unlock()
	if b.salt != nil {
		return b.salt, nil
	}
	salt, err := salt.NewSalt(b.view, &salt.Config{
		HashFunc: salt.SHA1Hash,
		Location: salt.DefaultLocation,
	})
	if err != nil {
		return nil, err
	}
	b.salt = salt
	b.MapAppId.SaltFunc = b.Salt
	b.MapUserId.SaltFunc = b.Salt

	return salt, nil
}

func (b *backend) invalidate(key string) {
	switch key {
	case salt.DefaultLocation:
		b.SaltMutex.Lock()
		defer b.SaltMutex.Unlock()
		b.salt = nil
	}
}

const backendHelp = `
The App ID credential provider is used to perform authentication from
within applications or machine by pairing together two hard-to-guess
unique pieces of information: a unique app ID, and a unique user ID.

The goal of this credential provider is to allow elastic users
(dynamic machines, containers, etc.) to authenticate with Vault without
having to store passwords outside of Vault. It is a single method of
solving the chicken-and-egg problem of setting up Vault access on a machine.
With this provider, nobody except the machine itself has access to both
pieces of information necessary to authenticate. For example:
configuration management will have the app IDs, but the machine itself
will detect its user ID based on some unique machine property such as a
MAC address (or a hash of it with some salt).

An example, real world process for using this provider:

  1. Create unique app IDs (UUIDs work well) and map them to policies.
     (Path: map/app-id/<app-id>)

  2. Store the app IDs within configuration management systems.

  3. An out-of-band process run by security operators map unique user IDs
     to these app IDs. Example: when an instance is launched, a cloud-init
     system tells security operators a unique ID for this machine. This
     process can be scripted, but the key is that it is out-of-band and
     out of reach of configuration management.
	 (Path: map/user-id/<user-id>)

  4. A new server is provisioned. Configuration management configures the
     app ID, the server itself detects its user ID. With both of these
     pieces of information, Vault can be accessed according to the policy
     set by the app ID.

More details on this process follow:

The app ID is a unique ID that maps to a set of policies. This ID is
generated by an operator and configured into the backend. The ID itself
is usually a UUID, but any hard-to-guess unique value can be used.

After creating app IDs, an operator authorizes a fixed set of user IDs
with each app ID. When a valid {app ID, user ID} tuple is given to the
"login" path, then the user is authenticated with the configured app
ID policies.

The user ID can be any value (just like the app ID), however it is
generally a value unique to a machine, such as a MAC address or instance ID,
or a value hashed from these unique values.

It is possible to authorize multiple app IDs with each
user ID by writing them as comma-separated values to the map/user-id/<user-id>
path.

It is also possible to renew the auth tokens with 'vault token-renew <token>' command.
Before the token is renewed, the validity of app ID, user ID and the associated
policies are checked again.
`
<commit_msg>Fix instantiation of salt funcs in app-id structs<commit_after>package appId

import (
	"sync"

	"github.com/hashicorp/vault/helper/salt"
	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/logical/framework"
)

func Factory(conf *logical.BackendConfig) (logical.Backend, error) {
	b, err := Backend(conf)
	if err != nil {
		return nil, err
	}
	return b.Setup(conf)
}

func Backend(conf *logical.BackendConfig) (*framework.Backend, error) {
	var b backend
	b.MapAppId = &framework.PolicyMap{
		PathMap: framework.PathMap{
			Name: "app-id",
			Schema: map[string]*framework.FieldSchema{
				"display_name": &framework.FieldSchema{
					Type:        framework.TypeString,
					Description: "A name to map to this app ID for logs.",
				},

				"value": &framework.FieldSchema{
					Type:        framework.TypeString,
					Description: "Policies for the app ID.",
				},
			},
		},
		DefaultKey: "default",
	}

	b.MapUserId = &framework.PathMap{
		Name: "user-id",
		Schema: map[string]*framework.FieldSchema{
			"cidr_block": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: "If not blank, restricts auth by this CIDR block",
			},

			"value": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: "App IDs that this user associates with.",
			},
		},
	}

	b.Backend = &framework.Backend{
		Help: backendHelp,

		PathsSpecial: &logical.Paths{
			Unauthenticated: []string{
				"login",
				"login/*",
			},
		},

		Paths: framework.PathAppend([]*framework.Path{
			pathLogin(&b),
			pathLoginWithAppIDPath(&b),
		},
			b.MapAppId.Paths(),
			b.MapUserId.Paths(),
		),

		AuthRenew: b.pathLoginRenew,

		Invalidate: b.invalidate,
	}

	b.view = conf.StorageView

	b.MapAppId.SaltFunc = b.Salt
	b.MapUserId.SaltFunc = b.Salt

	return b.Backend, nil
}

type backend struct {
	*framework.Backend

	salt      *salt.Salt
	SaltMutex sync.RWMutex
	view      logical.Storage
	MapAppId  *framework.PolicyMap
	MapUserId *framework.PathMap
}

func (b *backend) Salt() (*salt.Salt, error) {
	b.SaltMutex.RLock()
	if b.salt != nil {
		defer b.SaltMutex.RUnlock()
		return b.salt, nil
	}
	b.SaltMutex.RUnlock()
	b.SaltMutex.Lock()
	defer b.SaltMutex.Unlock()
	if b.salt != nil {
		return b.salt, nil
	}
	salt, err := salt.NewSalt(b.view, &salt.Config{
		HashFunc: salt.SHA1Hash,
		Location: salt.DefaultLocation,
	})
	if err != nil {
		return nil, err
	}
	b.salt = salt
	return salt, nil
}

func (b *backend) invalidate(key string) {
	switch key {
	case salt.DefaultLocation:
		b.SaltMutex.Lock()
		defer b.SaltMutex.Unlock()
		b.salt = nil
	}
}

const backendHelp = `
The App ID credential provider is used to perform authentication from
within applications or machine by pairing together two hard-to-guess
unique pieces of information: a unique app ID, and a unique user ID.

The goal of this credential provider is to allow elastic users
(dynamic machines, containers, etc.) to authenticate with Vault without
having to store passwords outside of Vault. It is a single method of
solving the chicken-and-egg problem of setting up Vault access on a machine.
With this provider, nobody except the machine itself has access to both
pieces of information necessary to authenticate. For example:
configuration management will have the app IDs, but the machine itself
will detect its user ID based on some unique machine property such as a
MAC address (or a hash of it with some salt).

An example, real world process for using this provider:

  1. Create unique app IDs (UUIDs work well) and map them to policies.
     (Path: map/app-id/<app-id>)

  2. Store the app IDs within configuration management systems.

  3. An out-of-band process run by security operators map unique user IDs
     to these app IDs. Example: when an instance is launched, a cloud-init
     system tells security operators a unique ID for this machine. This
     process can be scripted, but the key is that it is out-of-band and
     out of reach of configuration management.
	 (Path: map/user-id/<user-id>)

  4. A new server is provisioned. Configuration management configures the
     app ID, the server itself detects its user ID. With both of these
     pieces of information, Vault can be accessed according to the policy
     set by the app ID.

More details on this process follow:

The app ID is a unique ID that maps to a set of policies. This ID is
generated by an operator and configured into the backend. The ID itself
is usually a UUID, but any hard-to-guess unique value can be used.

After creating app IDs, an operator authorizes a fixed set of user IDs
with each app ID. When a valid {app ID, user ID} tuple is given to the
"login" path, then the user is authenticated with the configured app
ID policies.

The user ID can be any value (just like the app ID), however it is
generally a value unique to a machine, such as a MAC address or instance ID,
or a value hashed from these unique values.

It is possible to authorize multiple app IDs with each
user ID by writing them as comma-separated values to the map/user-id/<user-id>
path.

It is also possible to renew the auth tokens with 'vault token-renew <token>' command.
Before the token is renewed, the validity of app ID, user ID and the associated
policies are checked again.
`
<|endoftext|>