<commit_before>package smux

import (
	"container/heap"
	"encoding/binary"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
)

var (
	errInvalidProtocol = errors.New("invalid protocol")
	errGoAway          = errors.New("stream id overflows, should start a new connection")
	errTimeout         = errors.New("timeout")
)

type writeRequest struct {
	prio   uint64
	frame  Frame
	result chan writeResult
}

type writeResult struct {
	n   int
	err error
}

type buffersWriter interface {
	WriteBuffers(v [][]byte) (n int, err error)
}

// Session defines a multiplexed connection for streams
type Session struct {
	conn io.ReadWriteCloser

	config           *Config
	nextStreamID     uint32 // next stream identifier
	nextStreamIDLock sync.Mutex

	bucket       int32         // token bucket
	bucketNotify chan struct{} // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die     chan struct{} // flag session has died
	dieOnce sync.Once

	// socket error handling
	socketReadError      atomic.Value
	socketWriteError     atomic.Value
	chSocketReadError    chan struct{}
	chSocketWriteError   chan struct{}
	socketReadErrorOnce  sync.Once
	socketWriteErrorOnce sync.Once

	// smux protocol errors
	protoError     atomic.Value
	chProtoError   chan struct{}
	protoErrorOnce sync.Once

	chAccepts chan *Stream

	dataReady int32 // flag data has arrived

	goAway int32 // flag id exhausted

	deadline atomic.Value

	shaper chan writeRequest // a shaper for writing
	writes chan writeRequest
}

func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketNotify = make(chan struct{}, 1)
	s.shaper = make(chan writeRequest)
	s.writes = make(chan writeRequest)
	s.chSocketReadError = make(chan struct{})
	s.chSocketWriteError = make(chan struct{})
	s.chProtoError = make(chan struct{})

	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 0
	}

	go s.shaperLoop()
	go s.recvLoop()
	go s.sendLoop()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errors.WithStack(io.ErrClosedPipe)
	}

	// generate stream id
	s.nextStreamIDLock.Lock()
	if s.goAway > 0 {
		s.nextStreamIDLock.Unlock()
		return nil, errors.WithStack(errGoAway)
	}

	s.nextStreamID += 2
	sid := s.nextStreamID
	if sid == sid%2 { // stream-id overflows
		s.goAway = 1
		s.nextStreamIDLock.Unlock()
		return nil, errors.WithStack(errGoAway)
	}
	s.nextStreamIDLock.Unlock()

	stream := newStream(sid, s.config.MaxFrameSize, s)

	if _, err := s.writeFrame(newFrame(cmdSYN, sid)); err != nil {
		return nil, errors.WithStack(err)
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	select {
	case <-s.chSocketWriteError:
		return nil, s.socketWriteError.Load().(error)
	case <-s.die:
		return nil, errors.WithStack(io.ErrClosedPipe)
	default:
		s.streams[sid] = stream
		return stream, nil
	}
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	var deadline <-chan time.Time
	if d, ok := s.deadline.Load().(time.Time); ok && !d.IsZero() {
		timer := time.NewTimer(time.Until(d))
		defer timer.Stop()
		deadline = timer.C
	}

	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-deadline:
		return nil, errors.WithStack(errTimeout)
	case <-s.chSocketReadError:
		return nil, s.socketReadError.Load().(error)
	case <-s.chProtoError:
		return nil, s.protoError.Load().(error)
	case <-s.die:
		return nil, errors.WithStack(io.ErrClosedPipe)
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() error {
	var once bool
	s.dieOnce.Do(func() {
		close(s.die)
		once = true
	})

	if once {
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		return s.conn.Close()
	} else {
		return errors.WithStack(io.ErrClosedPipe)
	}
}

// notifyBucket notifies recvLoop that bucket is available
func (s *Session) notifyBucket() {
	select {
	case s.bucketNotify <- struct{}{}:
	default:
	}
}

func (s *Session) notifyReadError(err error) {
	s.socketReadErrorOnce.Do(func() {
		s.socketReadError.Store(err)
		close(s.chSocketReadError)
	})
}

func (s *Session) notifyWriteError(err error) {
	s.socketWriteErrorOnce.Do(func() {
		s.socketWriteError.Store(err)
		close(s.chSocketWriteError)
	})
}

func (s *Session) notifyProtoError(err error) {
	s.protoErrorOnce.Do(func() {
		s.protoError.Store(err)
		close(s.chProtoError)
	})
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// SetDeadline sets a deadline used by Accept* calls.
// A zero time value disables the deadline.
func (s *Session) SetDeadline(t time.Time) error {
	s.deadline.Store(t)
	return nil
}

// LocalAddr satisfies net.Conn interface
func (s *Session) LocalAddr() net.Addr {
	if ts, ok := s.conn.(interface {
		LocalAddr() net.Addr
	}); ok {
		return ts.LocalAddr()
	}
	return nil
}

// RemoteAddr satisfies net.Conn interface
func (s *Session) RemoteAddr() net.Addr {
	if ts, ok := s.conn.(interface {
		RemoteAddr() net.Addr
	}); ok {
		return ts.RemoteAddr()
	}
	return nil
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.notifyBucket()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.notifyBucket()
	}
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	var hdr rawHeader

	for {
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			select {
			case <-s.bucketNotify:
			case <-s.die:
				return
			}
		}

		// read header first
		if _, err := io.ReadFull(s.conn, hdr[:]); err == nil {
			atomic.StoreInt32(&s.dataReady, 1)
			if hdr.Version() != version {
				s.notifyProtoError(errors.WithStack(errInvalidProtocol))
				return
			}
			sid := hdr.StreamID()
			switch hdr.Cmd() {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[sid]; !ok {
					stream := newStream(sid, s.config.MaxFrameSize, s)
					s.streams[sid] = stream
					select {
					case s.chAccepts <- stream:
					case <-s.die:
					}
				}
				s.streamLock.Unlock()
			case cmdFIN:
				s.streamLock.Lock()
				if stream, ok := s.streams[sid]; ok {
					stream.fin()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				if hdr.Length() > 0 {
					newbuf := defaultAllocator.Get(int(hdr.Length()))
					if written, err := io.ReadFull(s.conn, newbuf); err == nil {
						s.streamLock.Lock()
						if stream, ok := s.streams[sid]; ok {
							stream.pushBytes(newbuf)
							atomic.AddInt32(&s.bucket, -int32(written))
							stream.notifyReadEvent()
						}
						s.streamLock.Unlock()
					} else {
						s.notifyReadError(errors.WithStack(err))
						return
					}
				}
			default:
				s.notifyProtoError(errors.WithStack(errInvalidProtocol))
				return
			}
		} else {
			s.notifyReadError(errors.WithStack(err))
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(s.config.KeepAliveInterval)
	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.writeFrameInternal(newFrame(cmdNOP, 0), tickerPing.C, 0)
			s.notifyBucket() // force a signal to the recvLoop
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

// shaper shapes the sending sequence among streams
func (s *Session) shaperLoop() {
	var reqs shaperHeap
	var next writeRequest
	var chWrite chan writeRequest

	for {
		if len(reqs) > 0 {
			chWrite = s.writes
			next = heap.Pop(&reqs).(writeRequest)
		} else {
			chWrite = nil
		}

		select {
		case <-s.die:
			return
		case r := <-s.shaper:
			if chWrite != nil { // next is valid, reshape
				heap.Push(&reqs, next)
			}
			heap.Push(&reqs, r)
		case chWrite <- next:
		}
	}
}

func (s *Session) sendLoop() {
	var buf []byte
	var n int
	var err error
	var vec [][]byte // vector for writeBuffers

	bw, ok := s.conn.(buffersWriter)
	if ok {
		buf = make([]byte, headerSize)
		vec = make([][]byte, 2)
	} else {
		buf = make([]byte, (1<<16)+headerSize)
	}

	for {
		select {
		case <-s.die:
			return
		case request := <-s.writes:
			buf[0] = request.frame.ver
			buf[1] = request.frame.cmd
			binary.LittleEndian.PutUint16(buf[2:], uint16(len(request.frame.data)))
			binary.LittleEndian.PutUint32(buf[4:], request.frame.sid)

			if len(vec) > 0 {
				vec[0] = buf[:headerSize]
				vec[1] = request.frame.data
				n, err = bw.WriteBuffers(vec)
			} else {
				copy(buf[headerSize:], request.frame.data)
				n, err = s.conn.Write(buf[:headerSize+len(request.frame.data)])
			}

			n -= headerSize
			if n < 0 {
				n = 0
			}

			result := writeResult{
				n:   n,
				err: errors.WithStack(err),
			}

			request.result <- result
			close(request.result)

			// store conn error
			if err != nil {
				s.notifyWriteError(errors.WithStack(err))
				return
			}
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(f Frame) (n int, err error) {
	return s.writeFrameInternal(f, nil, 0)
}

// internal writeFrame version to support deadline used in keepalive
func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time, prio uint64) (int, error) {
	req := writeRequest{
		prio:   prio,
		frame:  f,
		result: make(chan writeResult, 1),
	}
	select {
	case s.shaper <- req:
	case <-s.die:
		return 0, errors.WithStack(io.ErrClosedPipe)
	case <-s.chSocketWriteError:
		return 0, s.socketWriteError.Load().(error)
	case <-deadline:
		return 0, errors.WithStack(errTimeout)
	}

	select {
	case result := <-req.result:
		return result.n, errors.WithStack(result.err)
	case <-s.die:
		return 0, errors.WithStack(io.ErrClosedPipe)
	case <-s.chSocketWriteError:
		return 0, s.socketWriteError.Load().(error)
	case <-deadline:
		return 0, errors.WithStack(errTimeout)
	}
}
<commit_msg>add func Accept()<commit_after>package smux

import (
	"container/heap"
	"encoding/binary"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
)

var (
	errInvalidProtocol = errors.New("invalid protocol")
	errGoAway          = errors.New("stream id overflows, should start a new connection")
	errTimeout         = errors.New("timeout")
)

type writeRequest struct {
	prio   uint64
	frame  Frame
	result chan writeResult
}

type writeResult struct {
	n   int
	err error
}

type buffersWriter interface {
	WriteBuffers(v [][]byte) (n int, err error)
}

// Session defines a multiplexed connection for streams
type Session struct {
	conn io.ReadWriteCloser

	config           *Config
	nextStreamID     uint32 // next stream identifier
	nextStreamIDLock sync.Mutex

	bucket       int32         // token bucket
	bucketNotify chan struct{} // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die     chan struct{} // flag session has died
	dieOnce sync.Once

	// socket error handling
	socketReadError      atomic.Value
	socketWriteError     atomic.Value
	chSocketReadError    chan struct{}
	chSocketWriteError   chan struct{}
	socketReadErrorOnce  sync.Once
	socketWriteErrorOnce sync.Once

	// smux protocol errors
	protoError     atomic.Value
	chProtoError   chan struct{}
	protoErrorOnce sync.Once

	chAccepts chan *Stream

	dataReady int32 // flag data has arrived

	goAway int32 // flag id exhausted

	deadline atomic.Value

	shaper chan writeRequest // a shaper for writing
	writes chan writeRequest
}

func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketNotify = make(chan struct{}, 1)
	s.shaper = make(chan writeRequest)
	s.writes = make(chan writeRequest)
	s.chSocketReadError = make(chan struct{})
	s.chSocketWriteError = make(chan struct{})
	s.chProtoError = make(chan struct{})

	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 0
	}

	go s.shaperLoop()
	go s.recvLoop()
	go s.sendLoop()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errors.WithStack(io.ErrClosedPipe)
	}

	// generate stream id
	s.nextStreamIDLock.Lock()
	if s.goAway > 0 {
		s.nextStreamIDLock.Unlock()
		return nil, errors.WithStack(errGoAway)
	}

	s.nextStreamID += 2
	sid := s.nextStreamID
	if sid == sid%2 { // stream-id overflows
		s.goAway = 1
		s.nextStreamIDLock.Unlock()
		return nil, errors.WithStack(errGoAway)
	}
	s.nextStreamIDLock.Unlock()

	stream := newStream(sid, s.config.MaxFrameSize, s)

	if _, err := s.writeFrame(newFrame(cmdSYN, sid)); err != nil {
		return nil, errors.WithStack(err)
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	select {
	case <-s.chSocketWriteError:
		return nil, s.socketWriteError.Load().(error)
	case <-s.die:
		return nil, errors.WithStack(io.ErrClosedPipe)
	default:
		s.streams[sid] = stream
		return stream, nil
	}
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	var deadline <-chan time.Time
	if d, ok := s.deadline.Load().(time.Time); ok && !d.IsZero() {
		timer := time.NewTimer(time.Until(d))
		defer timer.Stop()
		deadline = timer.C
	}

	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-deadline:
		return nil, errors.WithStack(errTimeout)
	case <-s.chSocketReadError:
		return nil, s.socketReadError.Load().(error)
	case <-s.chProtoError:
		return nil, s.protoError.Load().(error)
	case <-s.die:
		return nil, errors.WithStack(io.ErrClosedPipe)
	}
}

// Accept Returns a generic ReadWriteCloser instead of smux.Stream
func (s *Session) Accept() (io.ReadWriteCloser, error) {
	return s.AcceptStream()
}

// Close is used to close the session and all streams.
func (s *Session) Close() error {
	var once bool
	s.dieOnce.Do(func() {
		close(s.die)
		once = true
	})

	if once {
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		return s.conn.Close()
	} else {
		return errors.WithStack(io.ErrClosedPipe)
	}
}

// notifyBucket notifies recvLoop that bucket is available
func (s *Session) notifyBucket() {
	select {
	case s.bucketNotify <- struct{}{}:
	default:
	}
}

func (s *Session) notifyReadError(err error) {
	s.socketReadErrorOnce.Do(func() {
		s.socketReadError.Store(err)
		close(s.chSocketReadError)
	})
}

func (s *Session) notifyWriteError(err error) {
	s.socketWriteErrorOnce.Do(func() {
		s.socketWriteError.Store(err)
		close(s.chSocketWriteError)
	})
}

func (s *Session) notifyProtoError(err error) {
	s.protoErrorOnce.Do(func() {
		s.protoError.Store(err)
		close(s.chProtoError)
	})
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// SetDeadline sets a deadline used by Accept* calls.
// A zero time value disables the deadline.
func (s *Session) SetDeadline(t time.Time) error {
	s.deadline.Store(t)
	return nil
}

// LocalAddr satisfies net.Conn interface
func (s *Session) LocalAddr() net.Addr {
	if ts, ok := s.conn.(interface {
		LocalAddr() net.Addr
	}); ok {
		return ts.LocalAddr()
	}
	return nil
}

// RemoteAddr satisfies net.Conn interface
func (s *Session) RemoteAddr() net.Addr {
	if ts, ok := s.conn.(interface {
		RemoteAddr() net.Addr
	}); ok {
		return ts.RemoteAddr()
	}
	return nil
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.notifyBucket()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.notifyBucket()
	}
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	var hdr rawHeader

	for {
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			select {
			case <-s.bucketNotify:
			case <-s.die:
				return
			}
		}

		// read header first
		if _, err := io.ReadFull(s.conn, hdr[:]); err == nil {
			atomic.StoreInt32(&s.dataReady, 1)
			if hdr.Version() != version {
				s.notifyProtoError(errors.WithStack(errInvalidProtocol))
				return
			}
			sid := hdr.StreamID()
			switch hdr.Cmd() {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[sid]; !ok {
					stream := newStream(sid, s.config.MaxFrameSize, s)
					s.streams[sid] = stream
					select {
					case s.chAccepts <- stream:
					case <-s.die:
					}
				}
				s.streamLock.Unlock()
			case cmdFIN:
				s.streamLock.Lock()
				if stream, ok := s.streams[sid]; ok {
					stream.fin()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				if hdr.Length() > 0 {
					newbuf := defaultAllocator.Get(int(hdr.Length()))
					if written, err := io.ReadFull(s.conn, newbuf); err == nil {
						s.streamLock.Lock()
						if stream, ok := s.streams[sid]; ok {
							stream.pushBytes(newbuf)
							atomic.AddInt32(&s.bucket, -int32(written))
							stream.notifyReadEvent()
						}
						s.streamLock.Unlock()
					} else {
						s.notifyReadError(errors.WithStack(err))
						return
					}
				}
			default:
				s.notifyProtoError(errors.WithStack(errInvalidProtocol))
				return
			}
		} else {
			s.notifyReadError(errors.WithStack(err))
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(s.config.KeepAliveInterval)
	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.writeFrameInternal(newFrame(cmdNOP, 0), tickerPing.C, 0)
			s.notifyBucket() // force a signal to the recvLoop
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

// shaper shapes the sending sequence among streams
func (s *Session) shaperLoop() {
	var reqs shaperHeap
	var next writeRequest
	var chWrite chan writeRequest

	for {
		if len(reqs) > 0 {
			chWrite = s.writes
			next = heap.Pop(&reqs).(writeRequest)
		} else {
			chWrite = nil
		}

		select {
		case <-s.die:
			return
		case r := <-s.shaper:
			if chWrite != nil { // next is valid, reshape
				heap.Push(&reqs, next)
			}
			heap.Push(&reqs, r)
		case chWrite <- next:
		}
	}
}

func (s *Session) sendLoop() {
	var buf []byte
	var n int
	var err error
	var vec [][]byte // vector for writeBuffers

	bw, ok := s.conn.(buffersWriter)
	if ok {
		buf = make([]byte, headerSize)
		vec = make([][]byte, 2)
	} else {
		buf = make([]byte, (1<<16)+headerSize)
	}

	for {
		select {
		case <-s.die:
			return
		case request := <-s.writes:
			buf[0] = request.frame.ver
			buf[1] = request.frame.cmd
			binary.LittleEndian.PutUint16(buf[2:], uint16(len(request.frame.data)))
			binary.LittleEndian.PutUint32(buf[4:], request.frame.sid)

			if len(vec) > 0 {
				vec[0] = buf[:headerSize]
				vec[1] = request.frame.data
				n, err = bw.WriteBuffers(vec)
			} else {
				copy(buf[headerSize:], request.frame.data)
				n, err = s.conn.Write(buf[:headerSize+len(request.frame.data)])
			}

			n -= headerSize
			if n < 0 {
				n = 0
			}

			result := writeResult{
				n:   n,
				err: errors.WithStack(err),
			}

			request.result <- result
			close(request.result)

			// store conn error
			if err != nil {
				s.notifyWriteError(errors.WithStack(err))
				return
			}
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(f Frame) (n int, err error) {
	return s.writeFrameInternal(f, nil, 0)
}

// internal writeFrame version to support deadline used in keepalive
func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time, prio uint64) (int, error) {
	req := writeRequest{
		prio:   prio,
		frame:  f,
		result: make(chan writeResult, 1),
	}
	select {
	case s.shaper <- req:
	case <-s.die:
		return 0, errors.WithStack(io.ErrClosedPipe)
	case <-s.chSocketWriteError:
		return 0, s.socketWriteError.Load().(error)
	case <-deadline:
		return 0, errors.WithStack(errTimeout)
	}

	select {
	case result := <-req.result:
		return result.n, errors.WithStack(result.err)
	case <-s.die:
		return 0, errors.WithStack(io.ErrClosedPipe)
	case <-s.chSocketWriteError:
		return 0, s.socketWriteError.Load().(error)
	case <-deadline:
		return 0, errors.WithStack(errTimeout)
	}
}
<|endoftext|>


<commit_before>package smux

import (
	"container/heap"
	"encoding/binary"
	"errors"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

const (
	defaultAcceptBacklog = 1024
)

var (
	ErrInvalidProtocol = errors.New("invalid protocol")
	ErrConsumed        = errors.New("peer consumed more than sent")
	ErrGoAway          = errors.New("stream id overflows, should start a new connection")
	ErrTimeout         = errors.New("timeout")
	ErrWouldBlock      = errors.New("operation would block on IO")
)

type writeRequest struct {
	prio   uint64
	frame  Frame
	result chan writeResult
}

type writeResult struct {
	n   int
	err error
}

type buffersWriter interface {
	WriteBuffers(v [][]byte) (n int, err error)
}

// Session defines a multiplexed connection for streams
type Session struct {
	conn io.ReadWriteCloser

	config           *Config
	nextStreamID     uint32 // next stream identifier
	nextStreamIDLock sync.Mutex

	bucket       int32         // token bucket
	bucketNotify chan struct{} // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die     chan struct{} // flag session has died
	dieOnce sync.Once

	// socket error handling
	socketReadError      atomic.Value
	socketWriteError     atomic.Value
	chSocketReadError    chan struct{}
	chSocketWriteError   chan struct{}
	socketReadErrorOnce  sync.Once
	socketWriteErrorOnce sync.Once

	// smux protocol errors
	protoError     atomic.Value
	chProtoError   chan struct{}
	protoErrorOnce sync.Once

	chAccepts chan *Stream

	dataReady int32 // flag data has arrived

	goAway int32 // flag id exhausted

	deadline atomic.Value

	shaper chan writeRequest // a shaper for writing
	writes chan writeRequest
}

func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketNotify = make(chan struct{}, 1)
	s.shaper = make(chan writeRequest)
	s.writes = make(chan writeRequest)
	s.chSocketReadError = make(chan struct{})
	s.chSocketWriteError = make(chan struct{})
	s.chProtoError = make(chan struct{})

	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 0
	}

	go s.shaperLoop()
	go s.recvLoop()
	go s.sendLoop()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, io.ErrClosedPipe
	}

	// generate stream id
	s.nextStreamIDLock.Lock()
	if s.goAway > 0 {
		s.nextStreamIDLock.Unlock()
		return nil, ErrGoAway
	}

	s.nextStreamID += 2
	sid := s.nextStreamID
	if sid == sid%2 { // stream-id overflows
		s.goAway = 1
		s.nextStreamIDLock.Unlock()
		return nil, ErrGoAway
	}
	s.nextStreamIDLock.Unlock()

	stream := newStream(sid, s.config.MaxFrameSize, s)

	if _, err := s.writeFrame(newFrame(byte(s.config.Version), cmdSYN, sid)); err != nil {
		return nil, err
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	select {
	case <-s.chSocketWriteError:
		return nil, s.socketWriteError.Load().(error)
	case <-s.die:
		return nil, io.ErrClosedPipe
	default:
		s.streams[sid] = stream
		return stream, nil
	}
}

// Open returns a generic ReadWriteCloser
func (s *Session) Open() (io.ReadWriteCloser, error) {
	return s.OpenStream()
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	var deadline <-chan time.Time
	if d, ok := s.deadline.Load().(time.Time); ok && !d.IsZero() {
		timer := time.NewTimer(time.Until(d))
		defer timer.Stop()
		deadline = timer.C
	}

	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-deadline:
		return nil, ErrTimeout
	case <-s.chSocketReadError:
		return nil, s.socketReadError.Load().(error)
	case <-s.chProtoError:
		return nil, s.protoError.Load().(error)
	case <-s.die:
		return nil, io.ErrClosedPipe
	}
}

// Accept Returns a generic ReadWriteCloser instead of smux.Stream
func (s *Session) Accept() (io.ReadWriteCloser, error) {
	return s.AcceptStream()
}

// Close is used to close the session and all streams.
func (s *Session) Close() error {
	var once bool
	s.dieOnce.Do(func() {
		close(s.die)
		once = true
	})

	if once {
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		return s.conn.Close()
	} else {
		return io.ErrClosedPipe
	}
}

// notifyBucket notifies recvLoop that bucket is available
func (s *Session) notifyBucket() {
	select {
	case s.bucketNotify <- struct{}{}:
	default:
	}
}

func (s *Session) notifyReadError(err error) {
	s.socketReadErrorOnce.Do(func() {
		s.socketReadError.Store(err)
		close(s.chSocketReadError)
	})
}

func (s *Session) notifyWriteError(err error) {
	s.socketWriteErrorOnce.Do(func() {
		s.socketWriteError.Store(err)
		close(s.chSocketWriteError)
	})
}

func (s *Session) notifyProtoError(err error) {
	s.protoErrorOnce.Do(func() {
		s.protoError.Store(err)
		close(s.chProtoError)
	})
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// SetDeadline sets a deadline used by Accept* calls.
// A zero time value disables the deadline.
func (s *Session) SetDeadline(t time.Time) error {
	s.deadline.Store(t)
	return nil
}

// LocalAddr satisfies net.Conn interface
func (s *Session) LocalAddr() net.Addr {
	if ts, ok := s.conn.(interface {
		LocalAddr() net.Addr
	}); ok {
		return ts.LocalAddr()
	}
	return nil
}

// RemoteAddr satisfies net.Conn interface
func (s *Session) RemoteAddr() net.Addr {
	if ts, ok := s.conn.(interface {
		RemoteAddr() net.Addr
	}); ok {
		return ts.RemoteAddr()
	}
	return nil
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.notifyBucket()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.notifyBucket()
	}
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	var hdr rawHeader
	var updHdr updHeader

	for {
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			select {
			case <-s.bucketNotify:
			case <-s.die:
				return
			}
		}

		// read header first
		if _, err := io.ReadFull(s.conn, hdr[:]); err == nil {
			atomic.StoreInt32(&s.dataReady, 1)
			if hdr.Version() != byte(s.config.Version) {
				s.notifyProtoError(ErrInvalidProtocol)
				return
			}
			sid := hdr.StreamID()
			switch hdr.Cmd() {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[sid]; !ok {
					stream := newStream(sid, s.config.MaxFrameSize, s)
					s.streams[sid] = stream
					select {
					case s.chAccepts <- stream:
					case <-s.die:
					}
				}
				s.streamLock.Unlock()
			case cmdFIN:
				s.streamLock.Lock()
				if stream, ok := s.streams[sid]; ok {
					stream.fin()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				if hdr.Length() > 0 {
					newbuf := defaultAllocator.Get(int(hdr.Length()))
					if written, err := io.ReadFull(s.conn, newbuf); err == nil {
						s.streamLock.Lock()
						if stream, ok := s.streams[sid]; ok {
							stream.pushBytes(newbuf)
							atomic.AddInt32(&s.bucket, -int32(written))
							stream.notifyReadEvent()
						}
						s.streamLock.Unlock()
					} else {
						s.notifyReadError(err)
						return
					}
				}
			case cmdUPD:
				if _, err := io.ReadFull(s.conn, updHdr[:]); err == nil {
					s.streamLock.Lock()
					if stream, ok := s.streams[sid]; ok {
						stream.update(updHdr.Consumed(), updHdr.Window())
					}
					s.streamLock.Unlock()
				} else {
					s.notifyReadError(err)
					return
				}
			default:
				s.notifyProtoError(ErrInvalidProtocol)
				return
			}
		} else {
			s.notifyReadError(err)
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(s.config.KeepAliveInterval)
	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.writeFrameInternal(newFrame(byte(s.config.Version), cmdNOP, 0), tickerPing.C, 0)
			s.notifyBucket() // force a signal to the recvLoop
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

// shaper shapes the sending sequence among streams
func (s *Session) shaperLoop() {
	var reqs shaperHeap
	var next writeRequest
	var chWrite chan writeRequest

	for {
		if len(reqs) > 0 {
			chWrite = s.writes
			next = heap.Pop(&reqs).(writeRequest)
		} else {
			chWrite = nil
		}

		select {
		case <-s.die:
			return
		case r := <-s.shaper:
			if chWrite != nil { // next is valid, reshape
				heap.Push(&reqs, next)
			}
			heap.Push(&reqs, r)
		case chWrite <- next:
		}
	}
}

func (s *Session) sendLoop() {
	var buf []byte
	var n int
	var err error
	var vec [][]byte // vector for writeBuffers

	bw, ok := s.conn.(buffersWriter)
	if ok {
		buf = make([]byte, headerSize)
		vec = make([][]byte, 2)
	} else {
		buf = make([]byte, (1<<16)+headerSize)
	}

	for {
		select {
		case <-s.die:
			return
		case request := <-s.writes:
			buf[0] = request.frame.ver
			buf[1] = request.frame.cmd
			binary.LittleEndian.PutUint16(buf[2:], uint16(len(request.frame.data)))
			binary.LittleEndian.PutUint32(buf[4:], request.frame.sid)

			if len(vec) > 0 {
				vec[0] = buf[:headerSize]
				vec[1] = request.frame.data
				n, err = bw.WriteBuffers(vec)
			} else {
				copy(buf[headerSize:], request.frame.data)
				n, err = s.conn.Write(buf[:headerSize+len(request.frame.data)])
			}

			n -= headerSize
			if n < 0 {
				n = 0
			}

			result := writeResult{
				n:   n,
				err: err,
			}

			request.result <- result
			close(request.result)

			// store conn error
			if err != nil {
				s.notifyWriteError(err)
				return
			}
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(f Frame) (n int, err error) {
	return s.writeFrameInternal(f, nil, 0)
}

// internal writeFrame version to support deadline used in keepalive
func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time, prio uint64) (int, error) {
	req := writeRequest{
		prio:   prio,
		frame:  f,
		result: make(chan writeResult, 1),
	}
	select {
	case s.shaper <- req:
	case <-s.die:
		return 0, io.ErrClosedPipe
	case <-s.chSocketWriteError:
		return 0, s.socketWriteError.Load().(error)
	case <-deadline:
		return 0, ErrTimeout
	}

	select {
	case result := <-req.result:
		return result.n, result.err
	case <-s.die:
		return 0, io.ErrClosedPipe
	case <-s.chSocketWriteError:
		return 0, s.socketWriteError.Load().(error)
	case <-deadline:
		return 0, ErrTimeout
	}
}
<commit_msg>OpenStream should return false if read has error issue#61<commit_after>package smux

import (
	"container/heap"
	"encoding/binary"
	"errors"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

const (
	defaultAcceptBacklog = 1024
)

var (
	ErrInvalidProtocol = errors.New("invalid protocol")
	ErrConsumed        = errors.New("peer consumed more than sent")
	ErrGoAway          = errors.New("stream id overflows, should start a new connection")
	ErrTimeout         = errors.New("timeout")
	ErrWouldBlock      = errors.New("operation would block on IO")
)

type writeRequest struct {
	prio   uint64
	frame  Frame
	result chan writeResult
}

type writeResult struct {
	n   int
	err error
}

type buffersWriter interface {
	WriteBuffers(v [][]byte) (n int, err error)
}

// Session defines a multiplexed connection for streams
type Session struct {
	conn io.ReadWriteCloser

	config           *Config
	nextStreamID     uint32 // next stream identifier
	nextStreamIDLock sync.Mutex

	bucket       int32         // token bucket
	bucketNotify chan struct{} // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die     chan struct{} // flag session has died
	dieOnce sync.Once

	// socket error handling
	socketReadError      atomic.Value
	socketWriteError     atomic.Value
	chSocketReadError    chan struct{}
	chSocketWriteError   chan struct{}
	socketReadErrorOnce  sync.Once
	socketWriteErrorOnce sync.Once

	// smux protocol errors
	protoError     atomic.Value
	chProtoError   chan struct{}
	protoErrorOnce sync.Once

	chAccepts chan *Stream

	dataReady int32 // flag data has arrived

	goAway int32 // flag id exhausted

	deadline atomic.Value

	shaper chan writeRequest // a shaper for writing
	writes chan writeRequest
}

func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketNotify = make(chan struct{}, 1)
	s.shaper = make(chan writeRequest)
	s.writes = make(chan writeRequest)
	s.chSocketReadError = make(chan struct{})
	s.chSocketWriteError = make(chan struct{})
	s.chProtoError = make(chan struct{})

	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 0
	}

	go s.shaperLoop()
	go s.recvLoop()
	go s.sendLoop()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, io.ErrClosedPipe
	}

	// generate stream id
	s.nextStreamIDLock.Lock()
	if s.goAway > 0 {
		s.nextStreamIDLock.Unlock()
		return nil, ErrGoAway
	}

	s.nextStreamID += 2
	sid := s.nextStreamID
	if sid == sid%2 { // stream-id overflows
		s.goAway = 1
		s.nextStreamIDLock.Unlock()
		return nil, ErrGoAway
	}
	s.nextStreamIDLock.Unlock()

	stream := newStream(sid, s.config.MaxFrameSize, s)

	if _, err := s.writeFrame(newFrame(byte(s.config.Version), cmdSYN, sid)); err != nil {
		return nil, err
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	select {
	case <-s.chSocketReadError:
		return nil, s.socketReadError.Load().(error)
	case <-s.chSocketWriteError:
		return nil, s.socketWriteError.Load().(error)
	case <-s.die:
		return nil, io.ErrClosedPipe
	default:
		s.streams[sid] = stream
		return stream, nil
	}
}

// Open returns a generic ReadWriteCloser
func (s *Session) Open() (io.ReadWriteCloser, error) {
	return s.OpenStream()
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	var deadline <-chan time.Time
	if d, ok := s.deadline.Load().(time.Time); ok && !d.IsZero() {
		timer := time.NewTimer(time.Until(d))
		defer timer.Stop()
		deadline = timer.C
	}

	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-deadline:
		return nil, ErrTimeout
	case <-s.chSocketReadError:
		return nil, s.socketReadError.Load().(error)
	case <-s.chProtoError:
		return nil, s.protoError.Load().(error)
	case <-s.die:
		return nil, io.ErrClosedPipe
	}
}

// Accept Returns a generic ReadWriteCloser instead of smux.Stream
func (s *Session) Accept() (io.ReadWriteCloser, error) {
	return s.AcceptStream()
}

// Close is used to close the session and all streams.
func (s *Session) Close() error {
	var once bool
	s.dieOnce.Do(func() {
		close(s.die)
		once = true
	})

	if once {
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		return s.conn.Close()
	} else {
		return io.ErrClosedPipe
	}
}

// notifyBucket notifies recvLoop that bucket is available
func (s *Session) notifyBucket() {
	select {
	case s.bucketNotify <- struct{}{}:
	default:
	}
}

func (s *Session) notifyReadError(err error) {
	s.socketReadErrorOnce.Do(func() {
		s.socketReadError.Store(err)
		close(s.chSocketReadError)
	})
}

func (s *Session) notifyWriteError(err error) {
	s.socketWriteErrorOnce.Do(func() {
		s.socketWriteError.Store(err)
		close(s.chSocketWriteError)
	})
}

func (s *Session) notifyProtoError(err error) {
	s.protoErrorOnce.Do(func() {
		s.protoError.Store(err)
		close(s.chProtoError)
	})
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// SetDeadline sets a deadline used by Accept* calls.
// A zero time value disables the deadline.
func (s *Session) SetDeadline(t time.Time) error {
	s.deadline.Store(t)
	return nil
}

// LocalAddr satisfies net.Conn interface
func (s *Session) LocalAddr() net.Addr {
	if ts, ok := s.conn.(interface {
		LocalAddr() net.Addr
	}); ok {
		return ts.LocalAddr()
	}
	return nil
}

// RemoteAddr satisfies net.Conn interface
func (s *Session) RemoteAddr() net.Addr {
	if ts, ok := s.conn.(interface {
		RemoteAddr() net.Addr
	}); ok {
		return ts.RemoteAddr()
	}
	return nil
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.notifyBucket()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.notifyBucket()
	}
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	var hdr rawHeader
	var updHdr updHeader

	for {
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			select {
			case <-s.bucketNotify:
			case <-s.die:
				return
			}
		}

		// read header first
		if _, err := io.ReadFull(s.conn, hdr[:]); err == nil {
			atomic.StoreInt32(&s.dataReady, 1)
			if hdr.Version() != byte(s.config.Version) {
				s.notifyProtoError(ErrInvalidProtocol)
				return
			}
			sid := hdr.StreamID()
			switch hdr.Cmd() {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[sid]; !ok {
					stream := newStream(sid, s.config.MaxFrameSize, s)
					s.streams[sid] = stream
					select {
					case s.chAccepts <- stream:
					case <-s.die:
					}
				}
				s.streamLock.Unlock()
			case cmdFIN:
				s.streamLock.Lock()
				if stream, ok := s.streams[sid]; ok {
					stream.fin()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				if hdr.Length() > 0 {
					newbuf := defaultAllocator.Get(int(hdr.Length()))
					if written, err := io.ReadFull(s.conn, newbuf); err == nil {
						s.streamLock.Lock()
						if stream, ok := s.streams[sid]; ok {
							stream.pushBytes(newbuf)
							atomic.AddInt32(&s.bucket, -int32(written))
							stream.notifyReadEvent()
						}
						s.streamLock.Unlock()
					} else {
						s.notifyReadError(err)
						return
					}
				}
			case cmdUPD:
				if _, err := io.ReadFull(s.conn, updHdr[:]); err == nil {
					s.streamLock.Lock()
					if stream, ok := s.streams[sid]; ok {
						stream.update(updHdr.Consumed(), updHdr.Window())
					}
					s.streamLock.Unlock()
				} else {
					s.notifyReadError(err)
					return
				}
			default:
				s.notifyProtoError(ErrInvalidProtocol)
				return
			}
		} else {
			s.notifyReadError(err)
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(s.config.KeepAliveInterval)
	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.writeFrameInternal(newFrame(byte(s.config.Version), cmdNOP, 0), tickerPing.C, 0)
			s.notifyBucket() // force a signal to the recvLoop
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

// shaper shapes the sending sequence among streams
func (s *Session) shaperLoop() {
	var reqs shaperHeap
	var next writeRequest
	var chWrite chan writeRequest

	for {
		if len(reqs) > 0 {
			chWrite = s.writes
			next = heap.Pop(&reqs).(writeRequest)
		} else {
			chWrite = nil
		}

		select {
		case <-s.die:
			return
		case r := <-s.shaper:
			if chWrite != nil { // next is valid, reshape
				heap.Push(&reqs, next)
			}
			heap.Push(&reqs, r)
		case chWrite <- next:
		}
	}
}

func (s *Session) sendLoop() {
	var buf []byte
	var n int
	var err error
	var vec [][]byte // vector for writeBuffers

	bw, ok := s.conn.(buffersWriter)
	if ok {
		buf = make([]byte, headerSize)
		vec = make([][]byte, 2)
	} else {
		buf = make([]byte, (1<<16)+headerSize)
	}

	for {
		select {
		case <-s.die:
			return
		case request := <-s.writes:
			buf[0] = request.frame.ver
			buf[1] = request.frame.cmd
			binary.LittleEndian.PutUint16(buf[2:], uint16(len(request.frame.data)))
			binary.LittleEndian.PutUint32(buf[4:], request.frame.sid)

			if len(vec) > 0 {
				vec[0] = buf[:headerSize]
				vec[1] = request.frame.data
				n, err = bw.WriteBuffers(vec)
			} else {
				copy(buf[headerSize:], request.frame.data)
				n, err = s.conn.Write(buf[:headerSize+len(request.frame.data)])
			}

			n -= headerSize
			if n < 0 {
				n = 0
			}

			result := writeResult{
				n:   n,
				err: err,
			}

			request.result <- result
			close(request.result)

			// store conn error
			if err != nil {
				s.notifyWriteError(err)
				return
			}
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(f Frame) (n int, err error) {
	return s.writeFrameInternal(f, nil, 0)
}

// internal writeFrame version to support deadline used in keepalive
func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time, prio uint64) (int, error) {
	req := writeRequest{
		prio:   prio,
		frame:  f,
		result: make(chan writeResult, 1),
	}
	select {
	case s.shaper <- req:
	case <-s.die:
		return 0, io.ErrClosedPipe
	case <-s.chSocketWriteError:
		return 0, s.socketWriteError.Load().(error)
	case <-deadline:
		return 0, ErrTimeout
	}

	select {
	case result := <-req.result:
		return result.n, result.err
	case <-s.die:
		return 0, io.ErrClosedPipe
	case <-s.chSocketWriteError:
		return 0, s.socketWriteError.Load().(error)
	case <-deadline:
		return 0, ErrTimeout
	}
}
<|endoftext|>


<commit_before>package signature

import (
	"testing"

	"github.com/containers/image/types"
	"github.com/docker/docker/reference"
)

// nameOnlyImageMock is a mock of types.Image which only allows transports.ImageName to work
type nameOnlyImageMock struct {
	forbiddenImageMock
}

func (nameOnlyImageMock) Reference() types.ImageReference {
	return nameOnlyImageReferenceMock("== StringWithinTransport mock")
}

// nameOnlyImageReferenceMock is a mock of types.ImageReference which only allows transports.ImageName to work, returning self.
type nameOnlyImageReferenceMock string

func (ref nameOnlyImageReferenceMock) Transport() types.ImageTransport {
	return nameImageTransportMock("== Transport mock")
}
func (ref nameOnlyImageReferenceMock) StringWithinTransport() string {
	return string(ref)
}
func (ref nameOnlyImageReferenceMock) DockerReference() reference.Named {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) PolicyConfigurationIdentity() string {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) PolicyConfigurationNamespaces() []string {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) NewImage(certPath string, tlsVerify bool) (types.Image, error) {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) NewImageSource(certPath string, tlsVerify bool) (types.ImageSource, error) {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) NewImageDestination(certPath string, tlsVerify bool) (types.ImageDestination, error) {
	panic("unexpected call to a mock function")
}

func TestPRInsecureAcceptAnythingIsSignatureAuthorAccepted(t *testing.T) {
	pr := NewPRInsecureAcceptAnything()
	// Pass nil signature to, kind of, test that the return value does not depend on it.
	sar, parsedSig, err := pr.isSignatureAuthorAccepted(nameOnlyImageMock{}, nil)
	assertSARUnknown(t, sar, parsedSig, err)
}

func TestPRInsecureAcceptAnythingIsRunningImageAllowed(t *testing.T) {
	pr := NewPRInsecureAcceptAnything()
	res, err := pr.isRunningImageAllowed(nameOnlyImageMock{})
	assertRunningAllowed(t, res, err)
}

func TestPRRejectIsSignatureAuthorAccepted(t *testing.T) {
	pr := NewPRReject()
	// Pass nil signature to, kind of, test that the return value does not depend on it.
	sar, parsedSig, err := pr.isSignatureAuthorAccepted(nameOnlyImageMock{}, nil)
	assertSARRejectedPolicyRequirement(t, sar, parsedSig, err)
}

func TestPRRejectIsRunningImageAllowed(t *testing.T) {
	// This will obviously need to change after this is implemented.
	pr := NewPRReject()
	res, err := pr.isRunningImageAllowed(nameOnlyImageMock{})
	assertRunningRejectedPolicyRequirement(t, res, err)
}
<commit_msg>Remove an obsolete comment<commit_after>package signature

import (
	"testing"

	"github.com/containers/image/types"
	"github.com/docker/docker/reference"
)

// nameOnlyImageMock is a mock of types.Image which only allows transports.ImageName to work
type nameOnlyImageMock struct {
	forbiddenImageMock
}

func (nameOnlyImageMock) Reference() types.ImageReference {
	return nameOnlyImageReferenceMock("== StringWithinTransport mock")
}

// nameOnlyImageReferenceMock is a mock of types.ImageReference which only allows transports.ImageName to work, returning self.
type nameOnlyImageReferenceMock string

func (ref nameOnlyImageReferenceMock) Transport() types.ImageTransport {
	return nameImageTransportMock("== Transport mock")
}
func (ref nameOnlyImageReferenceMock) StringWithinTransport() string {
	return string(ref)
}
func (ref nameOnlyImageReferenceMock) DockerReference() reference.Named {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) PolicyConfigurationIdentity() string {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) PolicyConfigurationNamespaces() []string {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) NewImage(certPath string, tlsVerify bool) (types.Image, error) {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) NewImageSource(certPath string, tlsVerify bool) (types.ImageSource, error) {
	panic("unexpected call to a mock function")
}
func (ref nameOnlyImageReferenceMock) NewImageDestination(certPath string, tlsVerify bool) (types.ImageDestination, error) {
	panic("unexpected call to a mock function")
}

func TestPRInsecureAcceptAnythingIsSignatureAuthorAccepted(t *testing.T) {
	pr := NewPRInsecureAcceptAnything()
	// Pass nil signature to, kind of, test that the return value does not depend on it.
	sar, parsedSig, err := pr.isSignatureAuthorAccepted(nameOnlyImageMock{}, nil)
	assertSARUnknown(t, sar, parsedSig, err)
}

func TestPRInsecureAcceptAnythingIsRunningImageAllowed(t *testing.T) {
	pr := NewPRInsecureAcceptAnything()
	res, err := pr.isRunningImageAllowed(nameOnlyImageMock{})
	assertRunningAllowed(t, res, err)
}

func TestPRRejectIsSignatureAuthorAccepted(t *testing.T) {
	pr := NewPRReject()
	// Pass nil signature to, kind of, test that the return value does not depend on it.
	sar, parsedSig, err := pr.isSignatureAuthorAccepted(nameOnlyImageMock{}, nil)
	assertSARRejectedPolicyRequirement(t, sar, parsedSig, err)
}

func TestPRRejectIsRunningImageAllowed(t *testing.T) {
	pr := NewPRReject()
	res, err := pr.isRunningImageAllowed(nameOnlyImageMock{})
	assertRunningRejectedPolicyRequirement(t, res, err)
}
<|endoftext|>


<commit_before>package gitmedia

import (
	core ".."
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
)

type ErrorsCommand struct {
	ClearLogs bool
	Boomtown  bool
	*Command
}

func (c *ErrorsCommand) Setup() {
	c.FlagSet.BoolVar(&c.ClearLogs, "clear", false, "Clear existing error logs")
	c.FlagSet.BoolVar(&c.Boomtown, "boomtown", false, "Trigger a panic")
}

func (c *ErrorsCommand) Run() {
	if c.ClearLogs {
		c.clear()
	}

	if c.Boomtown {
		c.boomtown()
		return
	}

	var sub string
	if len(c.SubCommands) > 0 {
		sub = c.SubCommands[0]
	}

	switch sub {
	case "last":
		c.lastError()
	case "":
		c.listErrors()
	default:
		core.Exit("Invalid errors sub command: %s", sub)
	}
}

func (c *ErrorsCommand) listErrors() {
	for _, path := range sortedLogs() {
		core.Print(path)
	}
}

func (c *ErrorsCommand) lastError() {
	logs := sortedLogs()
	last := logs[len(logs)-1]
	by, err := ioutil.ReadFile(filepath.Join(core.LocalLogDir, last))
	if err != nil {
		core.Panic(err, "Error reading log: %s", last)
	}

	core.Debug("Reading log: %s", last)
	os.Stdout.Write(by)
}

func (c *ErrorsCommand) clear() {
	err := os.RemoveAll(core.LocalLogDir)
	if err != nil {
		core.Panic(err, "Error clearing %s", core.LocalLogDir)
	}

	fmt.Println("Cleared", core.LocalLogDir)
}

func (c *ErrorsCommand) boomtown() {
	core.Debug("Debug message")
	err := errors.New("Error!")
	core.Panic(err, "Welcome to Boomtown")
	core.Debug("Never seen")
}

func sortedLogs() []string {
	fileinfos, err := ioutil.ReadDir(core.LocalLogDir)
	if err != nil {
		core.Panic(err, "Error reading logs directory: %s", core.LocalLogDir)
	}

	names := make([]string, len(fileinfos))
	for index, info := range fileinfos {
		names[index] = info.Name()
	}

	return names
}

func init() {
	registerCommand("errors", func(c *Command) RunnableCommand {
		return &ErrorsCommand{Command: c}
	})
}
<commit_msg>default errors subcommand reads the log file<commit_after>package gitmedia

import (
	core ".."
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
)

type ErrorsCommand struct {
	ClearLogs bool
	Boomtown  bool
	*Command
}

func (c *ErrorsCommand) Setup() {
	c.FlagSet.BoolVar(&c.ClearLogs, "clear", false, "Clear existing error logs")
	c.FlagSet.BoolVar(&c.Boomtown, "boomtown", false, "Trigger a panic")
}

func (c *ErrorsCommand) Run() {
	if c.ClearLogs {
		c.clear()
	}

	if c.Boomtown {
		c.boomtown()
		return
	}

	var sub string
	if len(c.SubCommands) > 0 {
		sub = c.SubCommands[0]
	}

	switch sub {
	case "last":
		c.lastError()
	case "":
		c.listErrors()
	default:
		c.showError(sub)
	}
}

func (c *ErrorsCommand) listErrors() {
	for _, path := range sortedLogs() {
		core.Print(path)
	}
}

func (c *ErrorsCommand) lastError() {
	logs := sortedLogs()
	c.showError(logs[len(logs)-1])
}

func (c *ErrorsCommand) showError(name string) {
	by, err := ioutil.ReadFile(filepath.Join(core.LocalLogDir, name))
	if err != nil {
		core.Panic(err, "Error reading log: %s", name)
	}

	core.Debug("Reading log: %s", name)
	os.Stdout.Write(by)
}

func (c *ErrorsCommand) clear() {
	err := os.RemoveAll(core.LocalLogDir)
	if err != nil {
		core.Panic(err, "Error clearing %s", core.LocalLogDir)
	}

	fmt.Println("Cleared", core.LocalLogDir)
}

func (c *ErrorsCommand) boomtown() {
	core.Debug("Debug message")
	err := errors.New("Error!")
	core.Panic(err, "Welcome to Boomtown")
	core.Debug("Never seen")
}

func sortedLogs() []string {
	fileinfos, err := ioutil.ReadDir(core.LocalLogDir)
	if err != nil {
		core.Panic(err, "Error reading logs directory: %s", core.LocalLogDir)
	}

	names := make([]string, len(fileinfos))
	for index, info := range fileinfos {
		names[index] = info.Name()
	}

	return names
}

func init() {
	registerCommand("errors", func(c *Command) RunnableCommand {
		return &ErrorsCommand{Command: c}
	})
}
<|endoftext|>


<commit_before>package magick

import (
	"bytes"
	"github.com/nfnt/resize"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

var (
	wizard = "test_data/wizard.png"
	newton = "test_data/Newtons_cradle_animation_book_2.gif"
	lenna  = "test_data/lenna.jpg"
)

func BenchmarkResizePng(b *testing.B) {
	im := decodeFile(b, "wizard.png")
	b.ResetTimer()
	for ii := 0; ii < b.N; ii++ {
		_, _ = im.Resize(240, 180, FLanczos)
	}
}

func BenchmarkResizePngGo(b *testing.B) {
	f, err := os.Open(wizard)
	if err != nil {
		b.Fatal(err)
	}
	defer f.Close()
	im, _, err := image.Decode(f)
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	for ii := 0; ii < b.N; ii++ {
		_ = resize.Resize(240, 180, im, resize.Lanczos2Lut)
	}
}

func benchmarkDecode(b *testing.B, file string) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	for ii := 0; ii < b.N; ii++ {
		if _, err := DecodeData(data); err != nil {
			b.Fatal(err)
		}
	}
}

func benchmarkDecodeGo(b *testing.B, file string) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		b.Fatal(err)
	}
	r := bytes.NewReader(data)
	if filepath.Ext(file) == ".gif" {
		b.ResetTimer()
		for ii := 0; ii < b.N; ii++ {
			if _, err := gif.DecodeAll(r); err != nil {
				b.Fatal(err)
			}
			r.Seek(0, 0)
		}
	} else {
		b.ResetTimer()
		for ii := 0; ii < b.N; ii++ {
			if _, _, err := image.Decode(r); err != nil {
				b.Fatal(err)
			}
			r.Seek(0, 0)
		}
	}
}

func benchmarkEncode(b *testing.B, file string, format string, quality uint) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		b.Fatal(err)
	}
	im, err := DecodeData(data)
	if err != nil {
		b.Fatal(err)
	}
	info := NewInfo()
	info.SetFormat(format)
	info.SetQuality(quality)
	b.ResetTimer()
	for ii := 0; ii < b.N; ii++ {
		if err := im.Encode(ioutil.Discard, info); err != nil {
			b.Fatal(err)
		}
	}
}

func benchmarkEncodeGo(b *testing.B, file string, format string, quality uint) {
	f, err := os.Open(file)
	if err != nil {
		b.Fatal(err)
	}
	defer f.Close()
	im, _, err := image.Decode(f)
	if err != nil {
		b.Fatal(err)
	}
	switch format {
	case "jpeg":
		opts := &jpeg.Options{
			Quality: int(quality),
		}
		b.ResetTimer()
		for ii := 0; ii < b.N; ii++ {
			if err := jpeg.Encode(ioutil.Discard, im, opts); err != nil {
				b.Fatal(err)
			}
		}
	case "png":
		b.ResetTimer()
		for ii := 0; ii < b.N; ii++ {
			if err := png.Encode(ioutil.Discard, im); err != nil {
				b.Fatal(err)
			}
		}
	default:
		b.Fatalf("format %s is not supported", format)
	}
}

func BenchmarkDecodePng(b *testing.B) {
	benchmarkDecode(b, wizard)
}

func BenchmarkDecodePngGo(b *testing.B) {
	benchmarkDecodeGo(b, wizard)
}

func BenchmarkDecodeGif(b *testing.B) {
	benchmarkDecode(b, newton)
}

func BenchmarkDecodeGifGo(b *testing.B) {
	benchmarkDecodeGo(b, newton)
}

func BenchmarkDecodeJpeg(b *testing.B) {
	benchmarkDecode(b, lenna)
}

func BenchmarkDecodeJpegGo(b *testing.B) {
	benchmarkDecodeGo(b, lenna)
}

func BenchmarkEncodePng(b *testing.B) {
	benchmarkEncode(b, lenna, "png", 0)
}

func BenchmarkEncodePngGo(b *testing.B) {
	benchmarkEncodeGo(b, lenna, "png", 0)
}

func BenchmarkEncodeJpeg(b *testing.B) {
	benchmarkEncode(b, wizard, "jpeg", 60)
}

func BenchmarkEncodeJpegGo(b *testing.B) {
	benchmarkEncodeGo(b, wizard, "jpeg", 60)
}

func BenchmarkEncodeJpegHQ(b *testing.B) {
	benchmarkEncode(b, wizard, "jpeg", 100)
}

func BenchmarkEncodeJpegHQGo(b *testing.B) {
	benchmarkEncodeGo(b, wizard, "jpeg", 100)
}

func BenchmarkEncodeJpegLQ(b *testing.B) {
	benchmarkEncode(b, wizard, "jpeg", 10)
}

func BenchmarkEncodeJpegLQGo(b *testing.B) {
	benchmarkEncodeGo(b, wizard, "jpeg", 10)
}
<commit_msg>Fix compile error due to changes in pkg github.com/nfnt/resize<commit_after>package magick

import (
	"bytes"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/nfnt/resize"
)

var (
	wizard = "test_data/wizard.png"
	newton = "test_data/Newtons_cradle_animation_book_2.gif"
	lenna  = "test_data/lenna.jpg"
)

func BenchmarkResizePng(b *testing.B) {
	im := decodeFile(b, "wizard.png")
	b.ResetTimer()
	for ii := 0; ii < b.N; ii++ {
		_, _ = im.Resize(240, 180, FLanczos)
	}
}

func BenchmarkResizePngGo(b *testing.B) {
	f, err := os.Open(wizard)
	if err != nil {
		b.Fatal(err)
	}
	defer f.Close()
	im, _, err := image.Decode(f)
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	for ii := 0; ii < b.N; ii++ {
		_ = resize.Resize(240, 180, im, resize.Lanczos2)
	}
}

func benchmarkDecode(b *testing.B, file string) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	for ii := 0; ii < b.N; ii++ {
		if _, err := DecodeData(data); err != nil {
			b.Fatal(err)
		}
	}
}

func benchmarkDecodeGo(b *testing.B, file string) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		b.Fatal(err)
	}
	r := bytes.NewReader(data)
	if filepath.Ext(file) == ".gif" {
		b.ResetTimer()
		for ii := 0; ii < b.N; ii++ {
			if _, err := gif.DecodeAll(r); err != nil {
				b.Fatal(err)
			}
			r.Seek(0, 0)
		}
	} else {
		b.ResetTimer()
		for ii := 0; ii < b.N; ii++ {
			if _, _, err := image.Decode(r); err != nil {
				b.Fatal(err)
			}
			r.Seek(0, 0)
		}
	}
}

func benchmarkEncode(b *testing.B, file string, format string, quality uint) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		b.Fatal(err)
	}
	im, err := DecodeData(data)
	if err != nil {
		b.Fatal(err)
	}
	info := NewInfo()
	info.SetFormat(format)
	info.SetQuality(quality)
	b.ResetTimer()
	for ii := 0; ii < b.N; ii++ {
		if err := im.Encode(ioutil.Discard, info); err != nil {
			b.Fatal(err)
		}
	}
}

func benchmarkEncodeGo(b *testing.B, file string, format string, quality uint) {
	f, err := os.Open(file)
	if err != nil {
		b.Fatal(err)
	}
	defer f.Close()
	im, _, err := image.Decode(f)
	if err != nil {
		b.Fatal(err)
	}
	switch format {
	case "jpeg":
		opts := &jpeg.Options{
			Quality: int(quality),
		}
		b.ResetTimer()
		for ii := 0; ii < b.N; ii++ {
			if err := jpeg.Encode(ioutil.Discard, im, opts); err != nil {
				b.Fatal(err)
			}
		}
	case "png":
		b.ResetTimer()
		for ii := 0; ii < b.N; ii++ {
			if err := png.Encode(ioutil.Discard, im); err != nil {
				b.Fatal(err)
			}
		}
	default:
		b.Fatalf("format %s is not supported", format)
	}
}

func BenchmarkDecodePng(b *testing.B) {
	benchmarkDecode(b, wizard)
}

func BenchmarkDecodePngGo(b *testing.B) {
	benchmarkDecodeGo(b, wizard)
}

func BenchmarkDecodeGif(b *testing.B) {
	benchmarkDecode(b, newton)
}

func BenchmarkDecodeGifGo(b *testing.B) {
	benchmarkDecodeGo(b, newton)
}

func BenchmarkDecodeJpeg(b *testing.B) {
	benchmarkDecode(b, lenna)
}

func BenchmarkDecodeJpegGo(b *testing.B) {
	benchmarkDecodeGo(b, lenna)
}

func BenchmarkEncodePng(b *testing.B) {
	benchmarkEncode(b, lenna, "png", 0)
}

func BenchmarkEncodePngGo(b *testing.B) {
	benchmarkEncodeGo(b, lenna, "png", 0)
}

func BenchmarkEncodeJpeg(b *testing.B) {
	benchmarkEncode(b, wizard, "jpeg", 60)
}

func BenchmarkEncodeJpegGo(b *testing.B) {
	benchmarkEncodeGo(b, wizard, "jpeg", 60)
}

func BenchmarkEncodeJpegHQ(b *testing.B) {
	benchmarkEncode(b, wizard, "jpeg", 100)
}

func BenchmarkEncodeJpegHQGo(b *testing.B) {
	benchmarkEncodeGo(b, wizard, "jpeg", 100)
}

func BenchmarkEncodeJpegLQ(b *testing.B) {
	benchmarkEncode(b, wizard, "jpeg", 10)
}

func BenchmarkEncodeJpegLQGo(b *testing.B) {
	benchmarkEncodeGo(b, wizard, "jpeg", 10)
}
<|endoftext|>


<commit_before>package eval

import "strings"

type alias struct {
	cmd  string
	args []string
}

var aliases = make(map[string]alias)

func DefAlias(name, arg string) {
	// TODO: Support more complex syntax.
	a := strings.Split(arg, " ")
	cmd := a[0]
	args := a[1:]
	if x, ok := aliases[cmd]; ok {
		cmd = x.cmd
		args = append(x.args, args...)
	}
	aliases[name] = alias{cmd, args}
}
<commit_msg>Strengthen the function of aliases a little<commit_after>package eval

import "strings"

type alias struct {
	cmd  string
	args []string
}

var aliases = make(map[string]alias)

func DefAlias(name, arg string) {
	// TODO: Support more complex syntax.
	a := strings.Split(arg, " ")
	cmd := a[0]
	args := a[1:]
	for i := range args {
		if args[i] == "''" {
			args[i] = ""
		}
	}
	if x, ok := aliases[cmd]; ok {
		cmd = x.cmd
		args = append(x.args, args...)
	}
	aliases[name] = alias{cmd, args}
}
<|endoftext|>


<commit_before>// a thin wrapper for common tasks involving os/signal
package signals

import (
	"os"
	ossignal "os/signal"
	"syscall"
)

const (
	SIGABRT   = syscall.SIGABRT
	SIGALRM   = syscall.SIGALRM
	SIGBUS    = syscall.SIGBUS
	SIGCHLD   = syscall.SIGCHLD
	SIGCONT   = syscall.SIGCONT
	SIGFPE    = syscall.SIGFPE
	SIGHUP    = syscall.SIGHUP
	SIGILL    = syscall.SIGILL
	SIGINT    = syscall.SIGINT
	SIGIO     = syscall.SIGIO
	SIGIOT    = syscall.SIGIOT
	SIGKILL   = syscall.SIGKILL
	SIGPIPE   = syscall.SIGPIPE
	SIGPROF   = syscall.SIGPROF
	SIGQUIT   = syscall.SIGQUIT
	SIGSEGV   = syscall.SIGSEGV
	SIGSTOP   = syscall.SIGSTOP
	SIGSYS    = syscall.SIGSYS
	SIGTERM   = syscall.SIGTERM
	SIGTRAP   = syscall.SIGTRAP
	SIGTSTP   = syscall.SIGTSTP
	SIGTTIN   = syscall.SIGTTIN
	SIGTTOU   = syscall.SIGTTOU
	SIGURG    = syscall.SIGURG
	SIGUSR1   = syscall.SIGUSR1
	SIGUSR2   = syscall.SIGUSR2
	SIGVTALRM = syscall.SIGVTALRM
	SIGWINCH  = syscall.SIGWINCH
	SIGXCPU   = syscall.SIGXCPU
	SIGXFSZ   = syscall.SIGXFSZ
)

type Handler func() bool

func Handle(signal os.Signal, handler Handler) {
	channel := make(chan os.Signal, 1)
	ossignal.Notify(channel, signal)
	for {
		<-channel
		if handler() == false {
			break
		}
	}
}
<commit_msg>don't leak channel<commit_after>// a thin wrapper for common tasks involving os/signal
package signals

import (
	"os"
	ossignal "os/signal"
	"syscall"
)

const (
	SIGABRT   = syscall.SIGABRT
	SIGALRM   = syscall.SIGALRM
	SIGBUS    = syscall.SIGBUS
	SIGCHLD   = syscall.SIGCHLD
	SIGCONT   = syscall.SIGCONT
	SIGFPE    = syscall.SIGFPE
	SIGHUP    = syscall.SIGHUP
	SIGILL    = syscall.SIGILL
	SIGINT    = syscall.SIGINT
	SIGIO     = syscall.SIGIO
	SIGIOT    = syscall.SIGIOT
	SIGKILL   = syscall.SIGKILL
	SIGPIPE   = syscall.SIGPIPE
	SIGPROF   = syscall.SIGPROF
	SIGQUIT   = syscall.SIGQUIT
	SIGSEGV   = syscall.SIGSEGV
	SIGSTOP   = syscall.SIGSTOP
	SIGSYS    = syscall.SIGSYS
	SIGTERM   = syscall.SIGTERM
	SIGTRAP   = syscall.SIGTRAP
	SIGTSTP   = syscall.SIGTSTP
	SIGTTIN   = syscall.SIGTTIN
	SIGTTOU   = syscall.SIGTTOU
	SIGURG    = syscall.SIGURG
	SIGUSR1   = syscall.SIGUSR1
	SIGUSR2   = syscall.SIGUSR2
	SIGVTALRM = syscall.SIGVTALRM
	SIGWINCH  = syscall.SIGWINCH
	SIGXCPU   = syscall.SIGXCPU
	SIGXFSZ   = syscall.SIGXFSZ
)

type Handler func() bool

func Handle(signal os.Signal, handler Handler) {
	channel := make(chan os.Signal, 1)
	ossignal.Notify(channel, signal)
	defer ossignal.Stop(c)

	for {
		<-channel
		if handler() == false {
			break
		}
	}
}
<|endoftext|>


<commit_before>package parser

import (
	"reflect"
	"strings"
	"testing"
)

var testVM = `
// comment
  push  constant	2
  // comment 2
push	 constant 3
	add 	// inline comment

push  constant	 1
sub	 //  inline comment 2
`

func TestHasMoreCommands(t *testing.T) {
	p := New(strings.NewReader(testVM))

	testCases := []struct {
		next   bool
		tokens []string
	}{
		{true, []string{"push", "constant", "2"}},
		{true, []string{"push", "constant", "3"}},
		{true, []string{"add"}},
		{true, []string{"push", "constant", "1"}},
		{true, []string{"sub"}},
		{false, []string{}},
	}

	for _, tt := range testCases {
		if p.HasMoreCommands() != tt.next {
			t.Errorf("HasMoreCommands should return %t, but %t", tt.next, !tt.next)
		}
		if !reflect.DeepEqual(p.tokens, tt.tokens) {
			t.Errorf("got %#v; want %#v", p.tokens, tt.tokens)
		}
	}
}

func TestAdvance(t *testing.T) {
	testCases := []struct {
		src  string
		want command
	}{
		{"add", command{Arithmetic, "add", 0}},
		{"sub", command{Arithmetic, "sub", 0}},
		{"push constant 1", command{Push, "constant", 1}},
		{"pop constant 2", command{Pop, "constant", 2}},
	}

	for _, tt := range testCases {
		p := New(strings.NewReader(tt.src))
		if p.HasMoreCommands() {
			if e := p.Advance(); e != nil {
				t.Errorf("advance failed: %s", e.Error())
			}
			if !reflect.DeepEqual(p.cmd, tt.want) {
				t.Errorf("got: %+v; want: %+v", p.cmd, tt.want)
			}
		}
	}
}

func TestAdvanceError(t *testing.T) {
	testCases := []struct {
		src string
	}{
		{"foo"},
		{"add sub"},
		{"push constant 1 2"},
		{"posh constant 1"},
		{"pop argment 0"},
		{"push local a"},
	}

	for _, tt := range testCases {
		p := New(strings.NewReader(tt.src))
		if p.HasMoreCommands() {
			if e := p.Advance(); e == nil {
				t.Errorf("expected error but got <nil>")
			}
		}
	}
}
<commit_msg>vmtranslator/parser: add a test case of label for HasMoreCommands<commit_after>package parser

import (
	"reflect"
	"strings"
	"testing"
)

var testVM = `
// comment
  push  constant	2
  // comment 2
push	 constant 3
	add 	// inline comment

push  constant	 1
sub	 //  inline comment 2

label LABEL0
`

func TestHasMoreCommands(t *testing.T) {
	p := New(strings.NewReader(testVM))

	testCases := []struct {
		next   bool
		tokens []string
	}{
		{true, []string{"push", "constant", "2"}},
		{true, []string{"push", "constant", "3"}},
		{true, []string{"add"}},
		{true, []string{"push", "constant", "1"}},
		{true, []string{"sub"}},
		{true, []string{"label", "LABEL0"}},
		{false, []string{}},
	}

	for _, tt := range testCases {
		if p.HasMoreCommands() != tt.next {
			t.Errorf("HasMoreCommands should return %t, but %t", tt.next, !tt.next)
		}
		if !reflect.DeepEqual(p.tokens, tt.tokens) {
			t.Errorf("got %#v; want %#v", p.tokens, tt.tokens)
		}
	}
}

func TestAdvance(t *testing.T) {
	testCases := []struct {
		src  string
		want command
	}{
		{"add", command{Arithmetic, "add", 0}},
		{"sub", command{Arithmetic, "sub", 0}},
		{"push constant 1", command{Push, "constant", 1}},
		{"pop constant 2", command{Pop, "constant", 2}},
	}

	for _, tt := range testCases {
		p := New(strings.NewReader(tt.src))
		if p.HasMoreCommands() {
			if e := p.Advance(); e != nil {
				t.Errorf("advance failed: %s", e.Error())
			}
			if !reflect.DeepEqual(p.cmd, tt.want) {
				t.Errorf("got: %+v; want: %+v", p.cmd, tt.want)
			}
		}
	}
}

func TestAdvanceError(t *testing.T) {
	testCases := []struct {
		src string
	}{
		{"foo"},
		{"add sub"},
		{"push constant 1 2"},
		{"posh constant 1"},
		{"pop argment 0"},
		{"push local a"},
	}

	for _, tt := range testCases {
		p := New(strings.NewReader(tt.src))
		if p.HasMoreCommands() {
			if e := p.Advance(); e == nil {
				t.Errorf("expected error but got <nil>")
			}
		}
	}
}
<|endoftext|>


<commit_before>package php

import "stephensearles.com/php/ast"

/*

Valid Expression Patterns
Expr [Binary Op] Expr
[Unary Op] Expr
Expr [Unary Op]
Expr [Tertiary Op 1] Expr [Tertiary Op 2] Expr
Identifier
Literal
Function Call

Parentesis always triggers sub-expression

non-associative clone new clone and new
left  [ array()
right ++ -- ~ (int) (float) (string) (array) (object) (bool) @  types and increment/decrement
non-associative instanceof  types
right ! logical
left  * / % arithmetic
left  + - . arithmetic and string
left  << >> bitwise
non-associative < <= > >= comparison
non-associative == != === !== <>  comparison
left  & bitwise and references
left  ^ bitwise
left  | bitwise
left  &&  logical
left  ||  logical
left  ? : ternary
right = += -= *= /= .= %= &= |= ^= <<= >>= => assignment
left  and logical
left  xor logical
left  or  logical
left  , many uses

*/

var operatorPrecedence = map[ItemType]int{
	itemArrayLookupOperatorLeft:  19,
	itemArrayLookupOperatorRight: 19,
	itemUnaryOperator:            18,
	itemCastOperator:             18,
	itemInstanceofOperator:       17,
	itemNegationOperator:         16,
	itemMultOperator:             15,
	itemAdditionOperator:         14,
	itemSubtractionOperator:      14,
	itemConcatenationOperator:    14,

	itemBitwiseShiftOperator: 13,
	itemComparisonOperator:   12,
	itemEqualityOperator:     11,

	itemAmpersandOperator:  10,
	itemBitwiseXorOperator: 9,
	itemBitwiseOrOperator:  8,
	itemAndOperator:        7,
	itemOrOperator:         6,
	itemTernaryOperator1:   5,
	itemTernaryOperator2:   5,
	itemAssignmentOperator: 4,
	itemWrittenAndOperator: 3,
	itemWrittenXorOperator: 2,
	itemWrittenOrOperator:  1,
}

func (p *parser) parseExpression() (expr ast.Expression) {
	// consume expression
	originalParenLev := p.parenLevel
	switch p.current.typ {
	case itemIgnoreErrorOperator:
		p.next()
		return p.parseExpression()
	case itemNewOperator:
		return &ast.NewExpression{
			Expression: p.parseNextExpression(),
		}
	case itemUnaryOperator, itemNegationOperator, itemAmpersandOperator, itemCastOperator:
		op := p.current
		expr = p.parseUnaryExpressionRight(p.parseNextExpression(), op)
	case itemArray:
		return p.parseArrayDeclaration()
	case itemIdentifier:
		if p.peek().typ == itemAssignmentOperator {
			assignee := p.parseIdentifier().(ast.Assignable)
			p.next()
			return ast.AssignmentExpression{
				Assignee: assignee,
				Operator: p.current.val,
				Value:    p.parseNextExpression(),
			}
		}
		fallthrough
	case itemNonVariableIdentifier, itemStringLiteral, itemNumberLiteral, itemBooleanLiteral, itemInclude, itemNull:
		expr = p.parseOperation(originalParenLev, p.expressionize())
	case itemOpenParen:
		p.parenLevel += 1
		p.next()
		expr = p.parseExpression()
		p.expect(itemCloseParen)
		p.parenLevel -= 1
		expr = p.parseOperation(originalParenLev, expr)
	default:
		p.errorf("Expected expression. Found %s", p.current)
		return
	}
	if p.parenLevel != originalParenLev {
		p.errorf("unbalanced parens: %d prev: %d", p.parenLevel, originalParenLev)
		return
	}
	return
}

func (p *parser) parseOperation(originalParenLevel int, lhs ast.Expression) (expr ast.Expression) {
	p.next()
	switch p.current.typ {
	case itemUnaryOperator:
		expr = p.parseUnaryExpressionLeft(lhs, p.current)
	case itemAdditionOperator, itemSubtractionOperator, itemConcatenationOperator, itemComparisonOperator, itemMultOperator, itemAndOperator, itemOrOperator, itemAmpersandOperator, itemBitwiseXorOperator, itemBitwiseOrOperator, itemBitwiseShiftOperator, itemWrittenAndOperator, itemWrittenXorOperator, itemWrittenOrOperator:
		expr = p.parseBinaryOperation(lhs, p.current, originalParenLevel)
	case itemTernaryOperator1:
		expr = p.parseTernaryOperation(lhs)
	case itemCloseParen:
		if p.parenLevel <= originalParenLevel {
			p.backup()
			return lhs
		}
		p.parenLevel -= 1
		return p.parseOperation(originalParenLevel, lhs)
	default:
		p.backup()
		return lhs
	}
	return p.parseOperation(originalParenLevel, expr)
}

func (p *parser) parseBinaryOperation(lhs ast.Expression, operator Item, originalParenLevel int) ast.Expression {
	p.next()
	rhs := p.expressionize()
	for {
		p.next()
		nextOperator := p.current
		p.backup()
		nextOperatorPrecedence, ok := operatorPrecedence[nextOperator.typ]
		if ok && nextOperatorPrecedence > operatorPrecedence[operator.typ] {
			rhs = p.parseOperation(originalParenLevel, rhs)
		} else {
			break
		}
	}
	return newBinaryOperation(operator, lhs, rhs)
}

func (p *parser) parseTernaryOperation(lhs ast.Expression) ast.Expression {
	truthy := p.parseNextExpression()
	p.expect(itemTernaryOperator2)
	falsy := p.parseNextExpression()
	return &ast.OperatorExpression{
		Operand1: lhs,
		Operand2: truthy,
		Operand3: falsy,
		Type:     truthy.EvaluatesTo() | falsy.EvaluatesTo(),
		Operator: "?:",
	}
}

func (p *parser) parseUnaryExpressionRight(operand ast.Expression, operator Item) ast.Expression {
	return newUnaryOperation(operator, operand)
}

func (p *parser) parseUnaryExpressionLeft(operand ast.Expression, operator Item) ast.Expression {
	return newUnaryOperation(operator, operand)
}

// expressionize takes the current token and returns it as the simplest
// expression for that token. That means an expression with no operators
// except for the object operator.
func (p *parser) expressionize() ast.Expression {
	switch p.current.typ {
	case itemIdentifier:
		return p.parseIdentifier()
	case itemStringLiteral, itemBooleanLiteral, itemNumberLiteral, itemNull:
		return p.parseLiteral()
	case itemNonVariableIdentifier:
		if p.peek().typ == itemOpenParen {
			var expr ast.Expression
			expr = p.parseFunctionCall()
			for p.peek().typ == itemObjectOperator {
				expr = p.parseObjectLookup(expr)
			}
			return expr
		}
		if p.peek().typ == itemScopeResolutionOperator {
			p.expect(itemScopeResolutionOperator)
			return &ast.ClassExpression{
				Receiver:   p.current.val,
				Expression: p.parseNextExpression(),
			}
		}
		return ast.ConstantExpression{
			Identifier: ast.NewIdentifier(p.current.val),
		}
	case itemOpenParen:
		return p.parseExpression()
	case itemInclude:
		return ast.Include{Expression: p.parseNextExpression()}
	}
	// error?
	return nil
}

func (p *parser) parseLiteral() *ast.Literal {
	switch p.current.typ {
	case itemStringLiteral:
		return &ast.Literal{Type: ast.String}
	case itemBooleanLiteral:
		return &ast.Literal{Type: ast.Boolean}
	case itemNumberLiteral:
		return &ast.Literal{Type: ast.Float}
	case itemNull:
		return &ast.Literal{Type: ast.Null}
	}
	p.errorf("Unknown literal type")
	return nil
}

func (p *parser) parseIdentifier() ast.Expression {
	var expr ast.Expression
	expr = ast.NewIdentifier(p.current.val)
	switch pk := p.peek(); pk.typ {
	case itemObjectOperator:
		for p.peek().typ == itemObjectOperator {
			expr = p.parseObjectLookup(expr)
		}
	case itemArrayLookupOperatorLeft:
		return p.parseArrayLookup(expr)
	}
	return expr
}

func (p *parser) parseObjectLookup(r ast.Expression) ast.Expression {
	p.expect(itemObjectOperator)
	p.expect(itemNonVariableIdentifier)
	switch pk := p.peek(); pk.typ {
	case itemOpenParen:
		expr := &ast.MethodCallExpression{
			Receiver:               r,
			FunctionCallExpression: p.parseFunctionCall(),
		}
		return expr
	case itemArrayLookupOperatorLeft:
		return p.parseArrayLookup(r)
	}
	return &ast.PropertyExpression{
		Receiver: r,
		Name:     p.current.val,
	}
}

func (p *parser) parseArrayLookup(e ast.Expression) ast.Expression {
	p.expect(itemArrayLookupOperatorLeft)
	if p.peek().typ == itemArrayLookupOperatorRight {
		p.expect(itemArrayLookupOperatorRight)
		return ast.ArrayAppendExpression{Array: e}
	}
	p.next()
	expr := &ast.ArrayLookupExpression{
		Array: e,
		Index: p.parseExpression(),
	}
	p.expect(itemArrayLookupOperatorRight)
	if p.peek().typ == itemArrayLookupOperatorLeft {
		return p.parseArrayLookup(expr)
	}
	return expr
}

func (p *parser) parseArrayDeclaration() ast.Expression {
	pairs := make([]ast.ArrayPair, 0)
	p.expect(itemOpenParen)
ArrayLoop:
	for {
		var key, val ast.Expression
		switch p.peek().typ {
		case itemCloseParen:
			break ArrayLoop
		default:
			val = p.parseNextExpression()
		}
		switch p.peek().typ {
		case itemComma:
			p.expect(itemComma)
		case itemCloseParen:
			pairs = append(pairs, ast.ArrayPair{Key: key, Value: val})
			break ArrayLoop
		case itemArrayKeyOperator:
			p.expect(itemArrayKeyOperator)
			key = val
			val = p.parseNextExpression()
			if p.peek().typ == itemCloseParen {
				pairs = append(pairs, ast.ArrayPair{Key: key, Value: val})
				break ArrayLoop
			}
			p.expect(itemComma)
		default:
			p.errorf("expected => or ,")
			return nil
		}
		pairs = append(pairs, ast.ArrayPair{Key: key, Value: val})
	}
	p.expect(itemCloseParen)
	return &ast.ArrayExpression{Pairs: pairs}
}
<commit_msg>Improved handling of complex expressions that begin with a negation operator<commit_after>package php

import "stephensearles.com/php/ast"

/*

Valid Expression Patterns
Expr [Binary Op] Expr
[Unary Op] Expr
Expr [Unary Op]
Expr [Tertiary Op 1] Expr [Tertiary Op 2] Expr
Identifier
Literal
Function Call

Parentesis always triggers sub-expression

non-associative clone new clone and new
left  [ array()
right ++ -- ~ (int) (float) (string) (array) (object) (bool) @  types and increment/decrement
non-associative instanceof  types
right ! logical
left  * / % arithmetic
left  + - . arithmetic and string
left  << >> bitwise
non-associative < <= > >= comparison
non-associative == != === !== <>  comparison
left  & bitwise and references
left  ^ bitwise
left  | bitwise
left  &&  logical
left  ||  logical
left  ? : ternary
right = += -= *= /= .= %= &= |= ^= <<= >>= => assignment
left  and logical
left  xor logical
left  or  logical
left  , many uses

*/

var operatorPrecedence = map[ItemType]int{
	itemArrayLookupOperatorLeft:  19,
	itemArrayLookupOperatorRight: 19,
	itemUnaryOperator:            18,
	itemCastOperator:             18,
	itemInstanceofOperator:       17,
	itemNegationOperator:         16,
	itemMultOperator:             15,
	itemAdditionOperator:         14,
	itemSubtractionOperator:      14,
	itemConcatenationOperator:    14,

	itemBitwiseShiftOperator: 13,
	itemComparisonOperator:   12,
	itemEqualityOperator:     11,

	itemAmpersandOperator:  10,
	itemBitwiseXorOperator: 9,
	itemBitwiseOrOperator:  8,
	itemAndOperator:        7,
	itemOrOperator:         6,
	itemTernaryOperator1:   5,
	itemTernaryOperator2:   5,
	itemAssignmentOperator: 4,
	itemWrittenAndOperator: 3,
	itemWrittenXorOperator: 2,
	itemWrittenOrOperator:  1,
}

func (p *parser) parseExpression() (expr ast.Expression) {
	// consume expression
	originalParenLev := p.parenLevel
	switch p.current.typ {
	case itemIgnoreErrorOperator:
		p.next()
		return p.parseExpression()
	case itemNewOperator:
		return &ast.NewExpression{
			Expression: p.parseNextExpression(),
		}
	case itemArray:
		return p.parseArrayDeclaration()
	case itemIdentifier:
		if p.peek().typ == itemAssignmentOperator {
			assignee := p.parseIdentifier().(ast.Assignable)
			p.next()
			return ast.AssignmentExpression{
				Assignee: assignee,
				Operator: p.current.val,
				Value:    p.parseNextExpression(),
			}
		}
		fallthrough
	case itemUnaryOperator, itemNegationOperator, itemAmpersandOperator, itemCastOperator:
		fallthrough
	case itemNonVariableIdentifier, itemStringLiteral, itemNumberLiteral, itemBooleanLiteral, itemInclude, itemNull:
		expr = p.parseOperation(originalParenLev, p.expressionize())
	case itemOpenParen:
		p.parenLevel += 1
		p.next()
		expr = p.parseExpression()
		p.expect(itemCloseParen)
		p.parenLevel -= 1
		expr = p.parseOperation(originalParenLev, expr)
	default:
		p.errorf("Expected expression. Found %s", p.current)
		return
	}
	if p.parenLevel != originalParenLev {
		p.errorf("unbalanced parens: %d prev: %d", p.parenLevel, originalParenLev)
		return
	}
	return
}

func (p *parser) parseOperation(originalParenLevel int, lhs ast.Expression) (expr ast.Expression) {
	p.next()
	switch p.current.typ {
	case itemUnaryOperator:
		expr = p.parseUnaryExpressionLeft(lhs, p.current)
	case itemAdditionOperator, itemSubtractionOperator, itemConcatenationOperator, itemComparisonOperator, itemMultOperator, itemAndOperator, itemOrOperator, itemAmpersandOperator, itemBitwiseXorOperator, itemBitwiseOrOperator, itemBitwiseShiftOperator, itemWrittenAndOperator, itemWrittenXorOperator, itemWrittenOrOperator:
		expr = p.parseBinaryOperation(lhs, p.current, originalParenLevel)
	case itemTernaryOperator1:
		expr = p.parseTernaryOperation(lhs)
	case itemCloseParen:
		if p.parenLevel <= originalParenLevel {
			p.backup()
			return lhs
		}
		p.parenLevel -= 1
		return p.parseOperation(originalParenLevel, lhs)
	case itemAssignmentOperator:
		assignee := lhs.(ast.Assignable)
		expr = ast.AssignmentExpression{
			Assignee: assignee,
			Operator: p.current.val,
			Value:    p.parseNextExpression(),
		}
	default:
		p.backup()
		return lhs
	}
	return p.parseOperation(originalParenLevel, expr)
}

func (p *parser) parseBinaryOperation(lhs ast.Expression, operator Item, originalParenLevel int) ast.Expression {
	p.next()
	rhs := p.expressionize()
	for {
		p.next()
		nextOperator := p.current
		p.backup()
		nextOperatorPrecedence, ok := operatorPrecedence[nextOperator.typ]
		if ok && nextOperatorPrecedence > operatorPrecedence[operator.typ] {
			rhs = p.parseOperation(originalParenLevel, rhs)
		} else {
			break
		}
	}
	return newBinaryOperation(operator, lhs, rhs)
}

func (p *parser) parseTernaryOperation(lhs ast.Expression) ast.Expression {
	truthy := p.parseNextExpression()
	p.expect(itemTernaryOperator2)
	falsy := p.parseNextExpression()
	return &ast.OperatorExpression{
		Operand1: lhs,
		Operand2: truthy,
		Operand3: falsy,
		Type:     truthy.EvaluatesTo() | falsy.EvaluatesTo(),
		Operator: "?:",
	}
}

func (p *parser) parseUnaryExpressionRight(operand ast.Expression, operator Item) ast.Expression {
	return newUnaryOperation(operator, operand)
}

func (p *parser) parseUnaryExpressionLeft(operand ast.Expression, operator Item) ast.Expression {
	return newUnaryOperation(operator, operand)
}

// expressionize takes the current token and returns it as the simplest
// expression for that token. That means an expression with no operators
// except for the object operator.
func (p *parser) expressionize() ast.Expression {
	switch p.current.typ {
	case itemUnaryOperator, itemNegationOperator:
		op := p.current
		p.next()
		return p.parseUnaryExpressionRight(p.expressionize(), op)
	case itemIdentifier:
		return p.parseIdentifier()
	case itemStringLiteral, itemBooleanLiteral, itemNumberLiteral, itemNull:
		return p.parseLiteral()
	case itemNonVariableIdentifier:
		if p.peek().typ == itemOpenParen {
			var expr ast.Expression
			expr = p.parseFunctionCall()
			for p.peek().typ == itemObjectOperator {
				expr = p.parseObjectLookup(expr)
			}
			return expr
		}
		if p.peek().typ == itemScopeResolutionOperator {
			p.expect(itemScopeResolutionOperator)
			return &ast.ClassExpression{
				Receiver:   p.current.val,
				Expression: p.parseNextExpression(),
			}
		}
		return ast.ConstantExpression{
			Identifier: ast.NewIdentifier(p.current.val),
		}
	case itemOpenParen:
		return p.parseExpression()
	case itemInclude:
		return ast.Include{Expression: p.parseNextExpression()}
	}
	// error?
	return nil
}

func (p *parser) parseLiteral() *ast.Literal {
	switch p.current.typ {
	case itemStringLiteral:
		return &ast.Literal{Type: ast.String}
	case itemBooleanLiteral:
		return &ast.Literal{Type: ast.Boolean}
	case itemNumberLiteral:
		return &ast.Literal{Type: ast.Float}
	case itemNull:
		return &ast.Literal{Type: ast.Null}
	}
	p.errorf("Unknown literal type")
	return nil
}

func (p *parser) parseIdentifier() ast.Expression {
	var expr ast.Expression
	expr = ast.NewIdentifier(p.current.val)
	switch pk := p.peek(); pk.typ {
	case itemObjectOperator:
		for p.peek().typ == itemObjectOperator {
			expr = p.parseObjectLookup(expr)
		}
	case itemArrayLookupOperatorLeft:
		return p.parseArrayLookup(expr)
	}
	return expr
}

func (p *parser) parseObjectLookup(r ast.Expression) ast.Expression {
	p.expect(itemObjectOperator)
	p.expect(itemNonVariableIdentifier)
	switch pk := p.peek(); pk.typ {
	case itemOpenParen:
		expr := &ast.MethodCallExpression{
			Receiver:               r,
			FunctionCallExpression: p.parseFunctionCall(),
		}
		return expr
	case itemArrayLookupOperatorLeft:
		return p.parseArrayLookup(r)
	}
	return &ast.PropertyExpression{
		Receiver: r,
		Name:     p.current.val,
	}
}

func (p *parser) parseArrayLookup(e ast.Expression) ast.Expression {
	p.expect(itemArrayLookupOperatorLeft)
	if p.peek().typ == itemArrayLookupOperatorRight {
		p.expect(itemArrayLookupOperatorRight)
		return ast.ArrayAppendExpression{Array: e}
	}
	p.next()
	expr := &ast.ArrayLookupExpression{
		Array: e,
		Index: p.parseExpression(),
	}
	p.expect(itemArrayLookupOperatorRight)
	if p.peek().typ == itemArrayLookupOperatorLeft {
		return p.parseArrayLookup(expr)
	}
	return expr
}

func (p *parser) parseArrayDeclaration() ast.Expression {
	pairs := make([]ast.ArrayPair, 0)
	p.expect(itemOpenParen)
ArrayLoop:
	for {
		var key, val ast.Expression
		switch p.peek().typ {
		case itemCloseParen:
			break ArrayLoop
		default:
			val = p.parseNextExpression()
		}
		switch p.peek().typ {
		case itemComma:
			p.expect(itemComma)
		case itemCloseParen:
			pairs = append(pairs, ast.ArrayPair{Key: key, Value: val})
			break ArrayLoop
		case itemArrayKeyOperator:
			p.expect(itemArrayKeyOperator)
			key = val
			val = p.parseNextExpression()
			if p.peek().typ == itemCloseParen {
				pairs = append(pairs, ast.ArrayPair{Key: key, Value: val})
				break ArrayLoop
			}
			p.expect(itemComma)
		default:
			p.errorf("expected => or ,")
			return nil
		}
		pairs = append(pairs, ast.ArrayPair{Key: key, Value: val})
	}
	p.expect(itemCloseParen)
	return &ast.ArrayExpression{Pairs: pairs}
}
<|endoftext|>