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

package ssh

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Listen requests the remote peer open a listening socket on
// addr. Incoming connections will be available by calling Accept on
// the returned net.Listener. The listener must be serviced, or the
// SSH connection may hang.
func (c *Client) Listen(n, addr string) (net.Listener, error) {
	laddr, err := net.ResolveTCPAddr(n, addr)
	if err != nil {
		return nil, err
	}
	return c.ListenTCP(laddr)
}

// Automatic port allocation is broken with OpenSSH before 6.0. See
// also https://bugzilla.mindrot.org/show_bug.cgi?id=2017.  In
// particular, OpenSSH 5.9 sends a channelOpenMsg with port number 0,
// rather than the actual port number. This means you can never open
// two different listeners with auto allocated ports. We work around
// this by trying explicit ports until we succeed.

const openSSHPrefix = "OpenSSH_"

var portRandomizer = rand.New(rand.NewSource(time.Now().UnixNano()))

// isBrokenOpenSSHVersion returns true if the given version string
// specifies a version of OpenSSH that is known to have a bug in port
// forwarding.
func isBrokenOpenSSHVersion(versionStr string) bool {
	i := strings.Index(versionStr, openSSHPrefix)
	if i < 0 {
		return false
	}
	i += len(openSSHPrefix)
	j := i
	for ; j < len(versionStr); j++ {
		if versionStr[j] < '0' || versionStr[j] > '9' {
			break
		}
	}
	version, _ := strconv.Atoi(versionStr[i:j])
	return version < 6
}

// autoPortListenWorkaround simulates automatic port allocation by
// trying random ports repeatedly.
func (c *Client) autoPortListenWorkaround(laddr *net.TCPAddr) (net.Listener, error) {
	var sshListener net.Listener
	var err error
	const tries = 10
	for i := 0; i < tries; i++ {
		addr := *laddr
		addr.Port = 1024 + portRandomizer.Intn(60000)
		sshListener, err = c.ListenTCP(&addr)
		if err == nil {
			laddr.Port = addr.Port
			return sshListener, err
		}
	}
	return nil, fmt.Errorf("ssh: listen on random port failed after %d tries: %v", tries, err)
}

// RFC 4254 7.1
type channelForwardMsg struct {
	addr  string
	rport uint32
}

// ListenTCP requests the remote peer open a listening socket
// on laddr. Incoming connections will be available by calling
// Accept on the returned net.Listener.
func (c *Client) ListenTCP(laddr *net.TCPAddr) (net.Listener, error) {
	if laddr.Port == 0 && isBrokenOpenSSHVersion(string(c.ServerVersion())) {
		return c.autoPortListenWorkaround(laddr)
	}

	m := channelForwardMsg{
		laddr.IP.String(),
		uint32(laddr.Port),
	}
	// send message
	ok, resp, err := c.SendRequest("tcpip-forward", true, Marshal(&m))
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, errors.New("ssh: tcpip-forward request denied by peer")
	}

	// If the original port was 0, then the remote side will
	// supply a real port number in the response.
	if laddr.Port == 0 {
		var p struct {
			Port uint32
		}
		if err := Unmarshal(resp, &p); err != nil {
			return nil, err
		}
		laddr.Port = int(p.Port)
	}

	// Register this forward, using the port number we obtained.
	ch := c.forwards.add(*laddr)

	return &tcpListener{laddr, c, ch}, nil
}

// forwardList stores a mapping between remote
// forward requests and the tcpListeners.
type forwardList struct {
	sync.Mutex
	entries []forwardEntry
}

// forwardEntry represents an established mapping of a laddr on a
// remote ssh server to a channel connected to a tcpListener.
type forwardEntry struct {
	laddr net.TCPAddr
	c     chan forward
}

// forward represents an incoming forwarded tcpip connection. The
// arguments to add/remove/lookup should be address as specified in
// the original forward-request.
type forward struct {
	newCh NewChannel   // the ssh client channel underlying this forward
	raddr *net.TCPAddr // the raddr of the incoming connection
}

func (l *forwardList) add(addr net.TCPAddr) chan forward {
	l.Lock()
	defer l.Unlock()
	f := forwardEntry{
		addr,
		make(chan forward, 1),
	}
	l.entries = append(l.entries, f)
	return f.c
}

// See RFC 4254, section 7.2
type forwardedTCPPayload struct {
	Addr       string
	Port       uint32
	OriginAddr string
	OriginPort uint32
}

// parseTCPAddr parses the originating address from the remote into a *net.TCPAddr.
func parseTCPAddr(addr string, port uint32) (*net.TCPAddr, error) {
	if port == 0 || port > 65535 {
		return nil, fmt.Errorf("ssh: port number out of range: %d", port)
	}
	ip := net.ParseIP(string(addr))
	if ip == nil {
		return nil, fmt.Errorf("ssh: cannot parse IP address %q", addr)
	}
	return &net.TCPAddr{IP: ip, Port: int(port)}, nil
}

func (l *forwardList) handleChannels(in <-chan NewChannel) {
	for ch := range in {
		var payload forwardedTCPPayload
		if err := Unmarshal(ch.ExtraData(), &payload); err != nil {
			ch.Reject(ConnectionFailed, "could not parse forwarded-tcpip payload: "+err.Error())
			continue
		}

		// RFC 4254 section 7.2 specifies that incoming
		// addresses should list the address, in string
		// format. It is implied that this should be an IP
		// address, as it would be impossible to connect to it
		// otherwise.
		laddr, err := parseTCPAddr(payload.Addr, payload.Port)
		if err != nil {
			ch.Reject(ConnectionFailed, err.Error())
			continue
		}
		raddr, err := parseTCPAddr(payload.OriginAddr, payload.OriginPort)
		if err != nil {
			ch.Reject(ConnectionFailed, err.Error())
			continue
		}

		if ok := l.forward(*laddr, *raddr, ch); !ok {
			// Section 7.2, implementations MUST reject spurious incoming
			// connections.
			ch.Reject(Prohibited, "no forward for address")
			continue
		}
	}
}

// remove removes the forward entry, and the channel feeding its
// listener.
func (l *forwardList) remove(addr net.TCPAddr) {
	l.Lock()
	defer l.Unlock()
	for i, f := range l.entries {
		if addr.IP.Equal(f.laddr.IP) && addr.Port == f.laddr.Port {
			l.entries = append(l.entries[:i], l.entries[i+1:]...)
			close(f.c)
			return
		}
	}
}

// closeAll closes and clears all forwards.
func (l *forwardList) closeAll() {
	l.Lock()
	defer l.Unlock()
	for _, f := range l.entries {
		close(f.c)
	}
	l.entries = nil
}

func (l *forwardList) forward(laddr, raddr net.TCPAddr, ch NewChannel) bool {
	l.Lock()
	defer l.Unlock()
	for _, f := range l.entries {
		if laddr.IP.Equal(f.laddr.IP) && laddr.Port == f.laddr.Port {
			f.c <- forward{ch, &raddr}
			return true
		}
	}
	return false
}

type tcpListener struct {
	laddr *net.TCPAddr

	conn *Client
	in   <-chan forward
}

// Accept waits for and returns the next connection to the listener.
func (l *tcpListener) Accept() (net.Conn, error) {
	s, ok := <-l.in
	if !ok {
		return nil, io.EOF
	}
	ch, incoming, err := s.newCh.Accept()
	if err != nil {
		return nil, err
	}
	go DiscardRequests(incoming)

	return &tcpChanConn{
		Channel: ch,
		laddr:   l.laddr,
		raddr:   s.raddr,
	}, nil
}

// Close closes the listener.
func (l *tcpListener) Close() error {
	m := channelForwardMsg{
		l.laddr.IP.String(),
		uint32(l.laddr.Port),
	}

	// this also closes the listener.
	l.conn.forwards.remove(*l.laddr)
	ok, _, err := l.conn.SendRequest("cancel-tcpip-forward", true, Marshal(&m))
	if err == nil && !ok {
		err = errors.New("ssh: cancel-tcpip-forward failed")
	}
	return err
}

// Addr returns the listener's network address.
func (l *tcpListener) Addr() net.Addr {
	return l.laddr
}

// Dial initiates a connection to the addr from the remote host.
// The resulting connection has a zero LocalAddr() and RemoteAddr().
func (c *Client) Dial(n, addr string) (net.Conn, error) {
	// Parse the address into host and numeric port.
	host, portString, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.ParseUint(portString, 10, 16)
	if err != nil {
		return nil, err
	}
	// Use a zero address for local and remote address.
	zeroAddr := &net.TCPAddr{
		IP:   net.IPv4zero,
		Port: 0,
	}
	ch, err := c.dial(net.IPv4zero.String(), 0, host, int(port))
	if err != nil {
		return nil, err
	}
	return &tcpChanConn{
		Channel: ch,
		laddr:   zeroAddr,
		raddr:   zeroAddr,
	}, nil
}

// DialTCP connects to the remote address raddr on the network net,
// which must be "tcp", "tcp4", or "tcp6".  If laddr is not nil, it is used
// as the local address for the connection.
func (c *Client) DialTCP(n string, laddr, raddr *net.TCPAddr) (net.Conn, error) {
	if laddr == nil {
		laddr = &net.TCPAddr{
			IP:   net.IPv4zero,
			Port: 0,
		}
	}
	ch, err := c.dial(laddr.IP.String(), laddr.Port, raddr.IP.String(), raddr.Port)
	if err != nil {
		return nil, err
	}
	return &tcpChanConn{
		Channel: ch,
		laddr:   laddr,
		raddr:   raddr,
	}, nil
}

// RFC 4254 7.2
type channelOpenDirectMsg struct {
	raddr string
	rport uint32
	laddr string
	lport uint32
}

func (c *Client) dial(laddr string, lport int, raddr string, rport int) (Channel, error) {
	msg := channelOpenDirectMsg{
		raddr: raddr,
		rport: uint32(rport),
		laddr: laddr,
		lport: uint32(lport),
	}
	ch, in, err := c.OpenChannel("direct-tcpip", Marshal(&msg))
	go DiscardRequests(in)
	return ch, err
}

type tcpChan struct {
	Channel // the backing channel
}

// tcpChanConn fulfills the net.Conn interface without
// the tcpChan having to hold laddr or raddr directly.
type tcpChanConn struct {
	Channel
	laddr, raddr net.Addr
}

// LocalAddr returns the local network address.
func (t *tcpChanConn) LocalAddr() net.Addr {
	return t.laddr
}

// RemoteAddr returns the remote network address.
func (t *tcpChanConn) RemoteAddr() net.Addr {
	return t.raddr
}

// SetDeadline sets the read and write deadlines associated
// with the connection.
func (t *tcpChanConn) SetDeadline(deadline time.Time) error {
	if err := t.SetReadDeadline(deadline); err != nil {
		return err
	}
	return t.SetWriteDeadline(deadline)
}

// SetReadDeadline sets the read deadline.
// A zero value for t means Read will not time out.
// After the deadline, the error from Read will implement net.Error
// with Timeout() == true.
func (t *tcpChanConn) SetReadDeadline(deadline time.Time) error {
	return errors.New("ssh: tcpChan: deadline not supported")
}

// SetWriteDeadline exists to satisfy the net.Conn interface
// but is not implemented by this type.  It always returns an error.
func (t *tcpChanConn) SetWriteDeadline(deadline time.Time) error {
	return errors.New("ssh: tcpChan: deadline not supported")
}
<commit_msg>crypto/ssh: Handle error in dial to avoid a goroutine leak<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssh

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Listen requests the remote peer open a listening socket on
// addr. Incoming connections will be available by calling Accept on
// the returned net.Listener. The listener must be serviced, or the
// SSH connection may hang.
func (c *Client) Listen(n, addr string) (net.Listener, error) {
	laddr, err := net.ResolveTCPAddr(n, addr)
	if err != nil {
		return nil, err
	}
	return c.ListenTCP(laddr)
}

// Automatic port allocation is broken with OpenSSH before 6.0. See
// also https://bugzilla.mindrot.org/show_bug.cgi?id=2017.  In
// particular, OpenSSH 5.9 sends a channelOpenMsg with port number 0,
// rather than the actual port number. This means you can never open
// two different listeners with auto allocated ports. We work around
// this by trying explicit ports until we succeed.

const openSSHPrefix = "OpenSSH_"

var portRandomizer = rand.New(rand.NewSource(time.Now().UnixNano()))

// isBrokenOpenSSHVersion returns true if the given version string
// specifies a version of OpenSSH that is known to have a bug in port
// forwarding.
func isBrokenOpenSSHVersion(versionStr string) bool {
	i := strings.Index(versionStr, openSSHPrefix)
	if i < 0 {
		return false
	}
	i += len(openSSHPrefix)
	j := i
	for ; j < len(versionStr); j++ {
		if versionStr[j] < '0' || versionStr[j] > '9' {
			break
		}
	}
	version, _ := strconv.Atoi(versionStr[i:j])
	return version < 6
}

// autoPortListenWorkaround simulates automatic port allocation by
// trying random ports repeatedly.
func (c *Client) autoPortListenWorkaround(laddr *net.TCPAddr) (net.Listener, error) {
	var sshListener net.Listener
	var err error
	const tries = 10
	for i := 0; i < tries; i++ {
		addr := *laddr
		addr.Port = 1024 + portRandomizer.Intn(60000)
		sshListener, err = c.ListenTCP(&addr)
		if err == nil {
			laddr.Port = addr.Port
			return sshListener, err
		}
	}
	return nil, fmt.Errorf("ssh: listen on random port failed after %d tries: %v", tries, err)
}

// RFC 4254 7.1
type channelForwardMsg struct {
	addr  string
	rport uint32
}

// ListenTCP requests the remote peer open a listening socket
// on laddr. Incoming connections will be available by calling
// Accept on the returned net.Listener.
func (c *Client) ListenTCP(laddr *net.TCPAddr) (net.Listener, error) {
	if laddr.Port == 0 && isBrokenOpenSSHVersion(string(c.ServerVersion())) {
		return c.autoPortListenWorkaround(laddr)
	}

	m := channelForwardMsg{
		laddr.IP.String(),
		uint32(laddr.Port),
	}
	// send message
	ok, resp, err := c.SendRequest("tcpip-forward", true, Marshal(&m))
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, errors.New("ssh: tcpip-forward request denied by peer")
	}

	// If the original port was 0, then the remote side will
	// supply a real port number in the response.
	if laddr.Port == 0 {
		var p struct {
			Port uint32
		}
		if err := Unmarshal(resp, &p); err != nil {
			return nil, err
		}
		laddr.Port = int(p.Port)
	}

	// Register this forward, using the port number we obtained.
	ch := c.forwards.add(*laddr)

	return &tcpListener{laddr, c, ch}, nil
}

// forwardList stores a mapping between remote
// forward requests and the tcpListeners.
type forwardList struct {
	sync.Mutex
	entries []forwardEntry
}

// forwardEntry represents an established mapping of a laddr on a
// remote ssh server to a channel connected to a tcpListener.
type forwardEntry struct {
	laddr net.TCPAddr
	c     chan forward
}

// forward represents an incoming forwarded tcpip connection. The
// arguments to add/remove/lookup should be address as specified in
// the original forward-request.
type forward struct {
	newCh NewChannel   // the ssh client channel underlying this forward
	raddr *net.TCPAddr // the raddr of the incoming connection
}

func (l *forwardList) add(addr net.TCPAddr) chan forward {
	l.Lock()
	defer l.Unlock()
	f := forwardEntry{
		addr,
		make(chan forward, 1),
	}
	l.entries = append(l.entries, f)
	return f.c
}

// See RFC 4254, section 7.2
type forwardedTCPPayload struct {
	Addr       string
	Port       uint32
	OriginAddr string
	OriginPort uint32
}

// parseTCPAddr parses the originating address from the remote into a *net.TCPAddr.
func parseTCPAddr(addr string, port uint32) (*net.TCPAddr, error) {
	if port == 0 || port > 65535 {
		return nil, fmt.Errorf("ssh: port number out of range: %d", port)
	}
	ip := net.ParseIP(string(addr))
	if ip == nil {
		return nil, fmt.Errorf("ssh: cannot parse IP address %q", addr)
	}
	return &net.TCPAddr{IP: ip, Port: int(port)}, nil
}

func (l *forwardList) handleChannels(in <-chan NewChannel) {
	for ch := range in {
		var payload forwardedTCPPayload
		if err := Unmarshal(ch.ExtraData(), &payload); err != nil {
			ch.Reject(ConnectionFailed, "could not parse forwarded-tcpip payload: "+err.Error())
			continue
		}

		// RFC 4254 section 7.2 specifies that incoming
		// addresses should list the address, in string
		// format. It is implied that this should be an IP
		// address, as it would be impossible to connect to it
		// otherwise.
		laddr, err := parseTCPAddr(payload.Addr, payload.Port)
		if err != nil {
			ch.Reject(ConnectionFailed, err.Error())
			continue
		}
		raddr, err := parseTCPAddr(payload.OriginAddr, payload.OriginPort)
		if err != nil {
			ch.Reject(ConnectionFailed, err.Error())
			continue
		}

		if ok := l.forward(*laddr, *raddr, ch); !ok {
			// Section 7.2, implementations MUST reject spurious incoming
			// connections.
			ch.Reject(Prohibited, "no forward for address")
			continue
		}
	}
}

// remove removes the forward entry, and the channel feeding its
// listener.
func (l *forwardList) remove(addr net.TCPAddr) {
	l.Lock()
	defer l.Unlock()
	for i, f := range l.entries {
		if addr.IP.Equal(f.laddr.IP) && addr.Port == f.laddr.Port {
			l.entries = append(l.entries[:i], l.entries[i+1:]...)
			close(f.c)
			return
		}
	}
}

// closeAll closes and clears all forwards.
func (l *forwardList) closeAll() {
	l.Lock()
	defer l.Unlock()
	for _, f := range l.entries {
		close(f.c)
	}
	l.entries = nil
}

func (l *forwardList) forward(laddr, raddr net.TCPAddr, ch NewChannel) bool {
	l.Lock()
	defer l.Unlock()
	for _, f := range l.entries {
		if laddr.IP.Equal(f.laddr.IP) && laddr.Port == f.laddr.Port {
			f.c <- forward{ch, &raddr}
			return true
		}
	}
	return false
}

type tcpListener struct {
	laddr *net.TCPAddr

	conn *Client
	in   <-chan forward
}

// Accept waits for and returns the next connection to the listener.
func (l *tcpListener) Accept() (net.Conn, error) {
	s, ok := <-l.in
	if !ok {
		return nil, io.EOF
	}
	ch, incoming, err := s.newCh.Accept()
	if err != nil {
		return nil, err
	}
	go DiscardRequests(incoming)

	return &tcpChanConn{
		Channel: ch,
		laddr:   l.laddr,
		raddr:   s.raddr,
	}, nil
}

// Close closes the listener.
func (l *tcpListener) Close() error {
	m := channelForwardMsg{
		l.laddr.IP.String(),
		uint32(l.laddr.Port),
	}

	// this also closes the listener.
	l.conn.forwards.remove(*l.laddr)
	ok, _, err := l.conn.SendRequest("cancel-tcpip-forward", true, Marshal(&m))
	if err == nil && !ok {
		err = errors.New("ssh: cancel-tcpip-forward failed")
	}
	return err
}

// Addr returns the listener's network address.
func (l *tcpListener) Addr() net.Addr {
	return l.laddr
}

// Dial initiates a connection to the addr from the remote host.
// The resulting connection has a zero LocalAddr() and RemoteAddr().
func (c *Client) Dial(n, addr string) (net.Conn, error) {
	// Parse the address into host and numeric port.
	host, portString, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.ParseUint(portString, 10, 16)
	if err != nil {
		return nil, err
	}
	// Use a zero address for local and remote address.
	zeroAddr := &net.TCPAddr{
		IP:   net.IPv4zero,
		Port: 0,
	}
	ch, err := c.dial(net.IPv4zero.String(), 0, host, int(port))
	if err != nil {
		return nil, err
	}
	return &tcpChanConn{
		Channel: ch,
		laddr:   zeroAddr,
		raddr:   zeroAddr,
	}, nil
}

// DialTCP connects to the remote address raddr on the network net,
// which must be "tcp", "tcp4", or "tcp6".  If laddr is not nil, it is used
// as the local address for the connection.
func (c *Client) DialTCP(n string, laddr, raddr *net.TCPAddr) (net.Conn, error) {
	if laddr == nil {
		laddr = &net.TCPAddr{
			IP:   net.IPv4zero,
			Port: 0,
		}
	}
	ch, err := c.dial(laddr.IP.String(), laddr.Port, raddr.IP.String(), raddr.Port)
	if err != nil {
		return nil, err
	}
	return &tcpChanConn{
		Channel: ch,
		laddr:   laddr,
		raddr:   raddr,
	}, nil
}

// RFC 4254 7.2
type channelOpenDirectMsg struct {
	raddr string
	rport uint32
	laddr string
	lport uint32
}

func (c *Client) dial(laddr string, lport int, raddr string, rport int) (Channel, error) {
	msg := channelOpenDirectMsg{
		raddr: raddr,
		rport: uint32(rport),
		laddr: laddr,
		lport: uint32(lport),
	}
	ch, in, err := c.OpenChannel("direct-tcpip", Marshal(&msg))
	if err != nil {
		return nil, err
	}
	go DiscardRequests(in)
	return ch, err
}

type tcpChan struct {
	Channel // the backing channel
}

// tcpChanConn fulfills the net.Conn interface without
// the tcpChan having to hold laddr or raddr directly.
type tcpChanConn struct {
	Channel
	laddr, raddr net.Addr
}

// LocalAddr returns the local network address.
func (t *tcpChanConn) LocalAddr() net.Addr {
	return t.laddr
}

// RemoteAddr returns the remote network address.
func (t *tcpChanConn) RemoteAddr() net.Addr {
	return t.raddr
}

// SetDeadline sets the read and write deadlines associated
// with the connection.
func (t *tcpChanConn) SetDeadline(deadline time.Time) error {
	if err := t.SetReadDeadline(deadline); err != nil {
		return err
	}
	return t.SetWriteDeadline(deadline)
}

// SetReadDeadline sets the read deadline.
// A zero value for t means Read will not time out.
// After the deadline, the error from Read will implement net.Error
// with Timeout() == true.
func (t *tcpChanConn) SetReadDeadline(deadline time.Time) error {
	return errors.New("ssh: tcpChan: deadline not supported")
}

// SetWriteDeadline exists to satisfy the net.Conn interface
// but is not implemented by this type.  It always returns an error.
func (t *tcpChanConn) SetWriteDeadline(deadline time.Time) error {
	return errors.New("ssh: tcpChan: deadline not supported")
}
<|endoftext|>


<commit_before>package main

import "os"
import "github.com/docker/engine-api/client"
import "math/rand"
import "os/exec"
import "strconv"
import "github.com/seehuhn/mt19937"
import "time"
import "os/user"

var cli *client.Client

type ExecRequest struct {
	Image string
	Cmd []string
	SourceFileName string
}

type Judge struct {
	Code string
	Compile *ExecRequest
	Exec ExecRequest
	Time int64
	Mem  int64
	TCCount int // The number of test cases
}

type JudgeResult int

const (
	Accepted            JudgeResult = 0
	WrongAnswer         JudgeResult = 1
	CompileError        JudgeResult = 2
	TimeLimitExceeded   JudgeResult = 3
	MemoryLimitExceeded JudgeResult = 4
	RuntimeError        JudgeResult = 5
	InternalError       JudgeResult = 6
	Judging             JudgeResult = 7
	CompileTimeLimitExceeded JudgeResult = 8
	CompileMemoryLimitExceeded JudgeResult = 9
)

type JudgeStatus struct {
	Case     *string
	JR       JudgeResult
	Mem      int64
	Time     int64
	Msg      *string
}

func CreateInternalError(msg string) JudgeStatus {
	return JudgeStatus{nil, InternalError, 0, 0, &msg}
}

const BASE_RAND_STRING = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

func RandomName() string {
	rng := rand.New(mt19937.New())
	rng.Seed(time.Now().UnixNano())

	res := make([]byte, 0, 32)
	for i := 0; i < 32; i++ {
		res = append(res, BASE_RAND_STRING[rng.Intn(len(BASE_RAND_STRING))])
	}

	return string(res)
}

func (j *Judge) Run(ch chan<- JudgeStatus, tests <-chan struct {
	Name string
	In  string
	Out string
}) {
	// Close a channel to send results of judging
	defer close(ch)

	// Identity
	id := RandomName()

	// User
	_, err := exec.Command("useradd", "--no-create-home", id).Output()

	if err != nil {
		ch <- CreateInternalError("Failed to create a directory to build your code. " + err.Error())

		return
	}

	uid, err := user.Lookup(id)

	if err != nil {
		ch <- CreateInternalError("Failed to look up a user. " + err.Error())

		return
	}

	uidInt, err := strconv.ParseInt(uid.Uid, 10, 64)

	if err != nil {
		ch <- CreateInternalError("Failed to parseInt uid. " + err.Error())

		return
	}
	gidInt, err := strconv.ParseInt(uid.Gid, 10, 64)
	if err != nil {
		ch <- CreateInternalError("Failed to parseInt gid. " + err.Error())

		return
	}

	defer exec.Command("userdel", id)

	// Working Directory
	path := workingDirectory + "/" + id

	err = os.Mkdir(path, 0664)

	if err != nil {
		ch <- CreateInternalError("Failed to create a directory. " + err.Error())

		return
	}

	defer os.RemoveAll(path)

	uidInt = uidInt * gidInt
	err = nil//os.Chown(path, int(uidInt), int(gidInt))

	if err != nil {
		ch <- CreateInternalError("Failed to chown the directory. " + err.Error())

		return
	}

	// Source File
	fp, err := os.Create(path + "/" + j.Compile.SourceFileName)

	if err != nil {
		ch <- CreateInternalError("Failed to create source file." + err.Error())

		return
	}

	l, err := fp.Write([]byte(j.Code))

	if err != nil {
		ch <- CreateInternalError("Failed to write your code on your file. " + err.Error())

		return
	}

	if l != len(j.Code) {
		ch <- CreateInternalError("Failed to write your code on your file.")

		return
	}

	fp.Close()

	err = os.Chmod(path + "/" + j.Compile.SourceFileName, 0644)

	if err != nil {
		ch <- CreateInternalError("Failed to chmod the source file. " + err.Error())

		return
	}

	// Compile
	if j.Compile != nil {
		exe, err := NewExecutor(id, 512 * 1024 * 1024, []string{"whoami"}, j.Compile.Image, []string{path + ":" + "/work"}, uid.Uid)

		if err != nil {
			ch <- CreateInternalError("Failed to create a Docker container to compile your code." + err.Error())

			return
		}

		res := exe.Run(10000, "")

		exe.Delete()
		if res.Status != ExecFinished {
			switch res.Status {
			case ExecError:
				ch <- CreateInternalError("Failed to execute a compiler." + res.Stderr)

				return
			case ExecMemoryLimitExceeded:
				ch <- JudgeStatus{JR: CompileMemoryLimitExceeded}

				return
			case ExecTimeLimitExceeded:
				ch <- JudgeStatus{JR: CompileTimeLimitExceeded}

				return
			}
		}

		if res.ExitCode != 0 || true { // Debug
			msg := res.Stdout + res.Stderr
			ch <- JudgeStatus{JR: CompileError, Msg: &msg}

			return
		}
	}
	return

	exe, err := NewExecutor(id, j.Mem, j.Exec.Cmd, j.Exec.Image, []string{path + ":" + "/work:ro"}, uid.Uid)

	if err != nil {
		ch <- CreateInternalError("Failed to create a Docker container to judge." + err.Error())

		return
	}

	defer exe.Delete()

	tcCounter := 0
	for tc, res := <-tests; res; tc, res = <-tests {
		res := exe.Run(j.Time, tc.In)

		if res.Status != ExecFinished {
			switch res.Status {
			case ExecError:
				msg := "Failed to execute your code." + res.Stderr
				ch <- JudgeStatus{Case: &tc.Name, JR: InternalError, Msg: &msg}
			case ExecMemoryLimitExceeded:
				ch <- JudgeStatus{Case: &tc.Name, JR: MemoryLimitExceeded}
			case ExecTimeLimitExceeded:
				ch <- JudgeStatus{Case: &tc.Name, JR: TimeLimitExceeded}
			}
		}else {
			if res.ExitCode != 0 {
				ch <- JudgeStatus{Case: &tc.Name, JR: RuntimeError}
			}else {
				if res.Stdout == tc.Out {
					ch <- JudgeStatus{Case: &tc.Name, JR: Accepted}
				}else {
					ch <- JudgeStatus{Case: &tc.Name, JR: WrongAnswer}
				}
			}
		}

		tcCounter++

		msg := strconv.FormatInt(int64(tcCounter), 10) + "/" + strconv.FormatInt(int64(j.TCCount), 10)
		ch <- JudgeStatus{JR: Judging, Msg: &msg}
	}

}
<commit_msg>Sun May 22 19:20:15 JST 2016<commit_after>package main

import "os"
import "github.com/docker/engine-api/client"
import "math/rand"
import "os/exec"
import "strconv"
import "github.com/seehuhn/mt19937"
import "time"
import "os/user"

var cli *client.Client

type ExecRequest struct {
	Image string
	Cmd []string
	SourceFileName string
}

type Judge struct {
	Code string
	Compile *ExecRequest
	Exec ExecRequest
	Time int64
	Mem  int64
	TCCount int // The number of test cases
}

type JudgeResult int

const (
	Accepted            JudgeResult = 0
	WrongAnswer         JudgeResult = 1
	CompileError        JudgeResult = 2
	TimeLimitExceeded   JudgeResult = 3
	MemoryLimitExceeded JudgeResult = 4
	RuntimeError        JudgeResult = 5
	InternalError       JudgeResult = 6
	Judging             JudgeResult = 7
	CompileTimeLimitExceeded JudgeResult = 8
	CompileMemoryLimitExceeded JudgeResult = 9
)

type JudgeStatus struct {
	Case     *string
	JR       JudgeResult
	Mem      int64
	Time     int64
	Msg      *string
}

func CreateInternalError(msg string) JudgeStatus {
	return JudgeStatus{nil, InternalError, 0, 0, &msg}
}

const BASE_RAND_STRING = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

func RandomName() string {
	rng := rand.New(mt19937.New())
	rng.Seed(time.Now().UnixNano())

	res := make([]byte, 0, 32)
	for i := 0; i < 32; i++ {
		res = append(res, BASE_RAND_STRING[rng.Intn(len(BASE_RAND_STRING))])
	}

	return string(res)
}

func (j *Judge) Run(ch chan<- JudgeStatus, tests <-chan struct {
	Name string
	In  string
	Out string
}) {
	// Close a channel to send results of judging
	defer close(ch)

	// Identity
	id := RandomName()

	// User
	_, err := exec.Command("useradd", "--no-create-home", id).Output()

	if err != nil {
		ch <- CreateInternalError("Failed to create a directory to build your code. " + err.Error())

		return
	}

	uid, err := user.Lookup(id)

	if err != nil {
		ch <- CreateInternalError("Failed to look up a user. " + err.Error())

		return
	}

	uidInt, err := strconv.ParseInt(uid.Uid, 10, 64)

	if err != nil {
		ch <- CreateInternalError("Failed to parseInt uid. " + err.Error())

		return
	}
	gidInt, err := strconv.ParseInt(uid.Gid, 10, 64)
	if err != nil {
		ch <- CreateInternalError("Failed to parseInt gid. " + err.Error())

		return
	}

	defer exec.Command("userdel", id)

	// Working Directory
	path := workingDirectory + "/" + id

	err = os.Mkdir(path, 0664)

	if err != nil {
		ch <- CreateInternalError("Failed to create a directory. " + err.Error())

		return
	}

	defer os.RemoveAll(path)

	uidInt = uidInt * gidInt
	err = nil//os.Chown(path, int(uidInt), int(gidInt))

	if err != nil {
		ch <- CreateInternalError("Failed to chown the directory. " + err.Error())

		return
	}

	// Source File
	fp, err := os.Create(path + "/" + j.Compile.SourceFileName)

	if err != nil {
		ch <- CreateInternalError("Failed to create source file." + err.Error())

		return
	}

	l, err := fp.Write([]byte(j.Code))

	if err != nil {
		ch <- CreateInternalError("Failed to write your code on your file. " + err.Error())

		return
	}

	if l != len(j.Code) {
		ch <- CreateInternalError("Failed to write your code on your file.")

		return
	}

	fp.Close()

	err = os.Chmod(path + "/" + j.Compile.SourceFileName, 0644)

	if err != nil {
		ch <- CreateInternalError("Failed to chmod the source file. " + err.Error())

		return
	}

	// Compile
	if j.Compile != nil {
		exe, err := NewExecutor(id, 512 * 1024 * 1024, []string{"ls", "-la", "/work"}, j.Compile.Image, []string{path + ":" + "/work"}, uid.Uid)

		if err != nil {
			ch <- CreateInternalError("Failed to create a Docker container to compile your code." + err.Error())

			return
		}

		res := exe.Run(10000, "")

		exe.Delete()
		if res.Status != ExecFinished {
			switch res.Status {
			case ExecError:
				ch <- CreateInternalError("Failed to execute a compiler." + res.Stderr)

				return
			case ExecMemoryLimitExceeded:
				ch <- JudgeStatus{JR: CompileMemoryLimitExceeded}

				return
			case ExecTimeLimitExceeded:
				ch <- JudgeStatus{JR: CompileTimeLimitExceeded}

				return
			}
		}

		if res.ExitCode != 0 || true { // Debug
			msg := res.Stdout + res.Stderr
			ch <- JudgeStatus{JR: CompileError, Msg: &msg}

			return
		}
	}
	return

	exe, err := NewExecutor(id, j.Mem, j.Exec.Cmd, j.Exec.Image, []string{path + ":" + "/work:ro"}, uid.Uid)

	if err != nil {
		ch <- CreateInternalError("Failed to create a Docker container to judge." + err.Error())

		return
	}

	defer exe.Delete()

	tcCounter := 0
	for tc, res := <-tests; res; tc, res = <-tests {
		res := exe.Run(j.Time, tc.In)

		if res.Status != ExecFinished {
			switch res.Status {
			case ExecError:
				msg := "Failed to execute your code." + res.Stderr
				ch <- JudgeStatus{Case: &tc.Name, JR: InternalError, Msg: &msg}
			case ExecMemoryLimitExceeded:
				ch <- JudgeStatus{Case: &tc.Name, JR: MemoryLimitExceeded}
			case ExecTimeLimitExceeded:
				ch <- JudgeStatus{Case: &tc.Name, JR: TimeLimitExceeded}
			}
		}else {
			if res.ExitCode != 0 {
				ch <- JudgeStatus{Case: &tc.Name, JR: RuntimeError}
			}else {
				if res.Stdout == tc.Out {
					ch <- JudgeStatus{Case: &tc.Name, JR: Accepted}
				}else {
					ch <- JudgeStatus{Case: &tc.Name, JR: WrongAnswer}
				}
			}
		}

		tcCounter++

		msg := strconv.FormatInt(int64(tcCounter), 10) + "/" + strconv.FormatInt(int64(j.TCCount), 10)
		ch <- JudgeStatus{JR: Judging, Msg: &msg}
	}

}
<|endoftext|>


<commit_before>package junos

import (
	"fmt"
	"github.com/Juniper/go-netconf/netconf"
	"log"
)

type Session struct {
	Conn *netconf.Session
}

func NewSession(host, user, password string) *Session {
    sess := &Session{}
    s, err := netconf.DialSSH(host, netconf.SSHConfigPassword(user, password))
    if err != nil {
        log.Fatal(err)
    }
    defer s.Close()
    sess.Conn = s

	return sess
}

func (s *Session) Lock() {
	resp, err := s.Conn.Exec("<rpc><lock><target><candidate/></target></lock></rpc>")

	if err != nil {
		fmt.Printf("Error: %+v\n", err)
	}

    fmt.Printf("%+v\n", resp)
}

func (s *Session) Unlock() {
	resp, err := s.Conn.Exec("<rpc><unlock><target><candidate/></target></unlock></rpc>")

	if err != nil {
		fmt.Printf("Error: %+v\n", err)
	}

    fmt.Printf("%+v\n", resp)
}<commit_msg>Changed RPC locking & unlocking<commit_after>package junos

import (
	"fmt"
	"github.com/Juniper/go-netconf/netconf"
	"log"
)

type Session struct {
	Conn *netconf.Session
}

func NewSession(host, user, password string) *Session {
    sess := &Session{}
    s, err := netconf.DialSSH(host, netconf.SSHConfigPassword(user, password))
    if err != nil {
        log.Fatal(err)
    }
    defer s.Close()
    sess.Conn = s

	return sess
}

func (s *Session) Lock() {
    fmt.Printf("%+v\n", s.Conn)
	// resp, err := s.Conn.Exec("<rpc><lock-configuration/></rpc>")

	// if err != nil {
		// fmt.Printf("Error: %+v\n", err)
	// }

    // fmt.Printf("%+v\n", resp)
}

func (s *Session) Unlock() {
    fmt.Printf("%+v\n", s.Conn)
	// resp, err := s.Conn.Exec("<rpc><unlock-configuration/></rpc>")

	// if err != nil {
		// fmt.Printf("Error: %+v\n", err)
	// }

    // fmt.Printf("%+v\n", resp)
}<|endoftext|>


<commit_before>package stats

import "sort"

type Map map[string]*Value

func (m Map) Values() Values {
	values := make(Values, 0, len(m))
	for _, v := range m {
		values = append(values, v)
	}
	return values
}

func (m Map) Inc(key string) int {
	return m.IncBy(key, 1)
}

func (m Map) IncBy(key string, value int) int {
	if m[key] == nil {
		m[key] = &Value{Key: key}
	}
	m[key].Value += value
	return m[key].Value
}

type Values []*Value

func (list Values) TopN(n int) Values {
	sort.Sort(sort.Reverse(list))
	out := make(Values, 0, n)
	for _, v := range list {
		out = append(out, v)
		if len(out) >= n {
			break
		}
	}
	return out
}

func (list Values) Len() int {
	return len(list)
}

func (list Values) Swap(a, b int) {
	list[a], list[b] = list[b], list[a]
}

func (list Values) Less(a, b int) bool {
	return list[a].Value < list[b].Value
}

type Value struct {
	Key   string
	Value int
}
<commit_msg>add convenience methods for SortedValues() and ReversedValues()<commit_after>package stats

import "sort"

type Map map[string]*Value

func (m Map) ReversedValues() Values {
	v := m.Values()
	sort.Sort(sort.Reverse(v))
	return v
}

func (m Map) SortedValues() Values {
	v := m.Values()
	sort.Sort(v)
	return v
}

func (m Map) Values() Values {
	values := make(Values, 0, len(m))
	for _, v := range m {
		values = append(values, v)
	}
	return values
}

func (m Map) Inc(key string) int {
	return m.IncBy(key, 1)
}

func (m Map) IncBy(key string, value int) int {
	if m[key] == nil {
		m[key] = &Value{Key: key}
	}
	m[key].Value += value
	return m[key].Value
}

type Values []*Value

func (list Values) TopN(n int) Values {
	sort.Sort(sort.Reverse(list))
	out := make(Values, 0, n)
	for _, v := range list {
		out = append(out, v)
		if len(out) >= n {
			break
		}
	}
	return out
}

func (list Values) Len() int {
	return len(list)
}

func (list Values) Swap(a, b int) {
	list[a], list[b] = list[b], list[a]
}

func (list Values) Less(a, b int) bool {
	return list[a].Value < list[b].Value
}

type Value struct {
	Key   string
	Value int
}
<|endoftext|>


<commit_before>package tui

import (
	"image"
	"strings"

	wordwrap "github.com/mitchellh/go-wordwrap"
)

var _ Widget = &Label{}

// Label is a widget to display read-only text.
type Label struct {
	WidgetBase

	text     string
	wordWrap bool

	// cache the result of SizeHint() (see #14)
	cacheSizeHint *image.Point

	styleName string
}

// NewLabel returns a new Label.
func NewLabel(text string) *Label {
	return &Label{
		text: text,
	}
}

// Draw draws the label.
func (l *Label) Draw(p *Painter) {
	lines := strings.Split(l.text, "\n")

	if l.wordWrap {
		lines = strings.Split(wordwrap.WrapString(l.text, uint(l.Size().X)), "\n")
	}

	style := "label"
	if l.styleName != "" {
		style += "." + l.styleName
	}

	p.WithStyle(style, func(p *Painter) {
		for i, line := range lines {
			p.DrawText(0, i, line)
		}
	})
}

// MinSizeHint returns the minimum size the widget is allowed to be.
func (l *Label) MinSizeHint() image.Point {
	return image.Point{1, 1}
}

// SizeHint returns the recommended size for the label.
func (l *Label) SizeHint() image.Point {
	if l.cacheSizeHint != nil {
		return *l.cacheSizeHint
	}
	var max int
	lines := strings.Split(l.text, "\n")
	for _, line := range lines {
		if w := stringWidth(line); w > max {
			max = w
		}
	}
	sizeHint := image.Point{max, l.heightForWidth(max)}
	l.cacheSizeHint = &sizeHint
	return sizeHint
}

func (l *Label) heightForWidth(w int) int {
	return len(strings.Split(wordwrap.WrapString(l.text, uint(w)), "\n"))
}

// SetText sets the text content of the label.
func (l *Label) SetText(text string) {
	l.cacheSizeHint = nil
	l.text = text
}

// SetWordWrap sets whether text content should be wrapped.
func (l *Label) SetWordWrap(enabled bool) {
	l.wordWrap = enabled
}

// SetStyleName sets the identifier used for custom styling.
func (l *Label) SetStyleName(style string) {
	l.styleName = style
}
<commit_msg>add text getter method and update SizeHint to account for word wrapping<commit_after>package tui

import (
	"image"
	"strings"

	wordwrap "github.com/mitchellh/go-wordwrap"
)

var _ Widget = &Label{}

// Label is a widget to display read-only text.
type Label struct {
	WidgetBase

	text     string
	wordWrap bool

	// cache the result of SizeHint() (see #14)
	cacheSizeHint *image.Point

	styleName string
}

// NewLabel returns a new Label.
func NewLabel(text string) *Label {
	return &Label{
		text: text,
	}
}

// Draw draws the label.
func (l *Label) Draw(p *Painter) {
	lines := strings.Split(l.text, "\n")

	if l.wordWrap {
		lines = strings.Split(wordwrap.WrapString(l.text, uint(l.Size().X)), "\n")
	}

	style := "label"
	if l.styleName != "" {
		style += "." + l.styleName
	}

	p.WithStyle(style, func(p *Painter) {
		for i, line := range lines {
			p.DrawText(0, i, line)
		}
	})
}

// MinSizeHint returns the minimum size the widget is allowed to be.
func (l *Label) MinSizeHint() image.Point {
	return image.Point{1, 1}
}

// SizeHint returns the recommended size for the label.
func (l *Label) SizeHint() image.Point {
	if l.cacheSizeHint != nil {
		return *l.cacheSizeHint
	}
	var max int
	lines := strings.Split(l.text, "\n")
	if l.wordWrap {
		lines = strings.Split(wordwrap.WrapString(l.text, uint(l.Size().X)), "\n")
	}
	for _, line := range lines {
		if w := stringWidth(line); w > max {
			max = w
		}
	}
	sizeHint := image.Point{max, l.heightForWidth(max)}
	l.cacheSizeHint = &sizeHint
	return sizeHint
}

func (l *Label) heightForWidth(w int) int {
	return len(strings.Split(wordwrap.WrapString(l.text, uint(w)), "\n"))
}

// GetText returns the text content of the label.
func (l *Label) GetText() string {
	return l.text
}

// SetText sets the text content of the label.
func (l *Label) SetText(text string) {
	l.cacheSizeHint = nil
	l.text = text
}

// SetWordWrap sets whether text content should be wrapped.
func (l *Label) SetWordWrap(enabled bool) {
	l.wordWrap = enabled
}

// SetStyleName sets the identifier used for custom styling.
func (l *Label) SetStyleName(style string) {
	l.styleName = style
}
<|endoftext|>


<commit_before>package jmespath

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"unicode/utf8"
)

type token struct {
	tokenType tokType
	value     string
	position  int
	length    int
}

type tokType int

const eof = -1

// Lexer contains information about the expression being tokenized.
type Lexer struct {
	expression string       // The expression provided by the user.
	currentPos int          // The current position in the string.
	lastWidth  int          // The width of the current rune.  This
	buf        bytes.Buffer // Internal buffer used for building up values.
}

// SyntaxError is the main error used whenever a lexing or parsing error occurs.
type SyntaxError struct {
	msg        string // Error message displayed to user
	Expression string // Expression that generated a SyntaxError
	Offset     int    // The location in the string where the error occurred
}

func (e SyntaxError) Error() string {
	// In the future, it would be good to underline the specific
	// location where the error occurred.
	return e.msg
}

//go:generate stringer -type=tokType
const (
	tUnknown tokType = iota
	tStar
	tDot
	tFilter
	tFlatten
	tLparen
	tRparen
	tLbracket
	tRbracket
	tLbrace
	tRbrace
	tOr
	tPipe
	tNumber
	tUnquotedIdentifier
	tQuotedIdentifier
	tComma
	tColon
	tLT
	tLTE
	tGT
	tGTE
	tEQ
	tNE
	tJSONLiteral
	tStringLiteral
	tCurrent
	tExpref
	tEOF
)

var basicTokens = map[rune]tokType{
	'.': tDot,
	'*': tStar,
	',': tComma,
	':': tColon,
	'{': tLbrace,
	'}': tRbrace,
	']': tRbracket, // tLbracket not included because it could be "[]"
	'(': tLparen,
	')': tRparen,
	'@': tCurrent,
	'&': tExpref,
}

var identifierStart = map[rune]bool{
	'a': true, 'b': true, 'c': true, 'd': true, 'e': true, 'f': true,
	'g': true, 'h': true, 'i': true, 'j': true, 'k': true, 'l': true,
	'm': true, 'n': true, 'o': true, 'p': true, 'q': true, 'r': true,
	's': true, 't': true, 'u': true, 'v': true, 'w': true, 'x': true,
	'y': true, 'z': true, 'A': true, 'B': true, 'C': true, 'D': true,
	'E': true, 'F': true, 'G': true, 'H': true, 'I': true, 'J': true,
	'K': true, 'L': true, 'M': true, 'N': true, 'O': true, 'P': true,
	'Q': true, 'R': true, 'S': true, 'T': true, 'U': true, 'V': true,
	'W': true, 'X': true, 'Y': true, 'Z': true, '_': true,
}

var identifierTrailing = map[rune]bool{
	'a': true, 'b': true, 'c': true, 'd': true, 'e': true, 'f': true,
	'g': true, 'h': true, 'i': true, 'j': true, 'k': true, 'l': true,
	'm': true, 'n': true, 'o': true, 'p': true, 'q': true, 'r': true,
	's': true, 't': true, 'u': true, 'v': true, 'w': true, 'x': true,
	'y': true, 'z': true, 'A': true, 'B': true, 'C': true, 'D': true,
	'E': true, 'F': true, 'G': true, 'H': true, 'I': true, 'J': true,
	'K': true, 'L': true, 'M': true, 'N': true, 'O': true, 'P': true,
	'Q': true, 'R': true, 'S': true, 'T': true, 'U': true, 'V': true,
	'W': true, 'X': true, 'Y': true, 'Z': true, '_': true, '0': true,
	'1': true, '2': true, '3': true, '4': true, '5': true, '6': true,
	'7': true, '8': true, '9': true,
}

var whiteSpace = map[rune]bool{
	' ': true, '\t': true, '\n': true, '\r': true,
}

func (t token) String() string {
	return fmt.Sprintf("Token{%+v, %s, %d, %d}",
		t.tokenType, t.value, t.position, t.length)
}

// NewLexer creates a new JMESPath lexer.
func NewLexer() *Lexer {
	lexer := Lexer{}
	return &lexer
}

func (lexer *Lexer) next() rune {
	if lexer.currentPos >= len(lexer.expression) {
		lexer.lastWidth = 0
		return eof
	}
	r, w := utf8.DecodeRuneInString(lexer.expression[lexer.currentPos:])
	lexer.lastWidth = w
	lexer.currentPos += w
	return r
}

func (lexer *Lexer) back() {
	lexer.currentPos -= lexer.lastWidth
}

func (lexer *Lexer) peek() rune {
	t := lexer.next()
	lexer.back()
	return t
}

// tokenize takes an expression and returns corresponding tokens.
func (lexer *Lexer) tokenize(expression string) ([]token, error) {
	var tokens []token
	lexer.expression = expression
	lexer.currentPos = 0
	lexer.lastWidth = 0
loop:
	for {
		r := lexer.next()
		if _, ok := identifierStart[r]; ok {
			t := lexer.consumeUnquotedIdentifier()
			tokens = append(tokens, t)
		} else if val, ok := basicTokens[r]; ok {
			// Basic single char token.
			t := token{
				tokenType: val,
				value:     string(r),
				position:  lexer.currentPos - lexer.lastWidth,
				length:    1,
			}
			tokens = append(tokens, t)
		} else if r == '-' || (r >= '0' && r <= '9') {
			t := lexer.consumeNumber()
			tokens = append(tokens, t)
		} else if r == '[' {
			t := lexer.consumeLBracket()
			tokens = append(tokens, t)
		} else if r == '"' {
			t, err := lexer.consumeQuotedIdentifier()
			if err != nil {
				return tokens, err
			}
			tokens = append(tokens, t)
		} else if r == '\'' {
			t, err := lexer.consumeRawStringLiteral()
			if err != nil {
				return tokens, err
			}
			tokens = append(tokens, t)
		} else if r == '`' {
			t, err := lexer.consumeLiteral()
			if err != nil {
				return tokens, err
			}
			tokens = append(tokens, t)
		} else if r == '|' {
			t := lexer.matchOrElse(r, '|', tOr, tPipe)
			tokens = append(tokens, t)
		} else if r == '<' {
			t := lexer.matchOrElse(r, '=', tLTE, tLT)
			tokens = append(tokens, t)
		} else if r == '>' {
			t := lexer.matchOrElse(r, '=', tGTE, tGT)
			tokens = append(tokens, t)
		} else if r == '!' {
			t := lexer.matchOrElse(r, '=', tNE, tUnknown)
			tokens = append(tokens, t)
		} else if r == '=' {
			t := lexer.matchOrElse(r, '=', tEQ, tUnknown)
			tokens = append(tokens, t)
		} else if r == eof {
			break loop
		} else if _, ok := whiteSpace[r]; ok {
			// Ignore whitespace
		} else {
			return tokens, lexer.syntaxError(fmt.Sprintf("Unknown char: %s", strconv.QuoteRuneToASCII(r)))
		}
	}
	tokens = append(tokens, token{tEOF, "", len(lexer.expression), 0})
	return tokens, nil
}

// Consume characters until the ending rune "r" is reached.
// If the end of the expression is reached before seeing the
// terminating rune "r", then an error is returned.
// If no error occurs then the matching substring is returned.
// The returned string will not include the ending rune.
func (lexer *Lexer) consumeUntil(end rune) (string, error) {
	start := lexer.currentPos
	current := lexer.next()
	for current != end && current != eof {
		if current == '\\' && lexer.peek() != eof {
			lexer.next()
		}
		current = lexer.next()
	}
	if lexer.lastWidth == 0 {
		// Then we hit an EOF so we never reached the closing
		// delimiter.
		return "", &SyntaxError{
			msg:        "Unclosed delimiter: " + string(end),
			Expression: lexer.expression,
			Offset:     len(lexer.expression),
		}
	}
	return lexer.expression[start : lexer.currentPos-lexer.lastWidth], nil
}

func (lexer *Lexer) consumeLiteral() (token, error) {
	start := lexer.currentPos
	value, err := lexer.consumeUntil('`')
	if err != nil {
		return token{}, err
	}
	value = strings.Replace(value, "\\`", "`", -1)
	return token{
		tokenType: tJSONLiteral,
		value:     value,
		position:  start,
		length:    len(value),
	}, nil
}

func (lexer *Lexer) consumeRawStringLiteral() (token, error) {
	start := lexer.currentPos
	currentIndex := start
	current := lexer.next()
	for current != '\'' && lexer.peek() != eof {
		if current == '\\' && lexer.peek() == '\'' {
			chunk := lexer.expression[currentIndex : lexer.currentPos-1]
			lexer.buf.WriteString(chunk)
			lexer.buf.WriteString("'")
			lexer.next()
			currentIndex = lexer.currentPos
		}
		current = lexer.next()
	}
	if lexer.lastWidth == 0 {
		// Then we hit an EOF so we never reached the closing
		// delimiter.
		return token{}, &SyntaxError{
			msg:        "Unclosed delimiter: '",
			Expression: lexer.expression,
			Offset:     len(lexer.expression),
		}
	}
	if currentIndex < lexer.currentPos {
		lexer.buf.WriteString(lexer.expression[currentIndex : lexer.currentPos-1])
	}
	value := lexer.buf.String()
	// Reset the buffer so it can reused again.
	lexer.buf.Reset()
	return token{
		tokenType: tStringLiteral,
		value:     value,
		position:  start,
		length:    len(value),
	}, nil
}

func (lexer *Lexer) syntaxError(msg string) SyntaxError {
	return SyntaxError{
		msg:        msg,
		Expression: lexer.expression,
		Offset:     lexer.currentPos,
	}
}

// Checks for a two char token, otherwise matches a single character
// token. This is used whenever a two char token overlaps a single
// char token, e.g. "||" -> tPipe, "|" -> tOr.
func (lexer *Lexer) matchOrElse(first rune, second rune, matchedType tokType, singleCharType tokType) token {
	start := lexer.currentPos - lexer.lastWidth
	nextRune := lexer.next()
	var t token
	if nextRune == second {
		t = token{
			tokenType: matchedType,
			value:     string(first) + string(second),
			position:  start,
			length:    2,
		}
	} else {
		lexer.back()
		t = token{
			tokenType: singleCharType,
			value:     string(first),
			position:  start,
			length:    1,
		}
	}
	return t
}

func (lexer *Lexer) consumeLBracket() token {
	// There's three options here:
	// 1. A filter expression "[?"
	// 2. A flatten operator "[]"
	// 3. A bare rbracket "["
	start := lexer.currentPos - lexer.lastWidth
	nextRune := lexer.next()
	var t token
	if nextRune == '?' {
		t = token{
			tokenType: tFilter,
			value:     "[?",
			position:  start,
			length:    2,
		}
	} else if nextRune == ']' {
		t = token{
			tokenType: tFlatten,
			value:     "[]",
			position:  start,
			length:    2,
		}
	} else {
		t = token{
			tokenType: tLbracket,
			value:     "[",
			position:  start,
			length:    1,
		}
		lexer.back()
	}
	return t
}

func (lexer *Lexer) consumeQuotedIdentifier() (token, error) {
	start := lexer.currentPos
	value, err := lexer.consumeUntil('"')
	if err != nil {
		return token{}, err
	}
	var decoded string
	asJSON := []byte("\"" + value + "\"")
	if err := json.Unmarshal([]byte(asJSON), &decoded); err != nil {
		return token{}, err
	}
	return token{
		tokenType: tQuotedIdentifier,
		value:     decoded,
		position:  start - 1,
		length:    len(decoded),
	}, nil
}

func (lexer *Lexer) consumeUnquotedIdentifier() token {
	// Consume runes until we reach the end of an unquoted
	// identifier.
	start := lexer.currentPos - lexer.lastWidth
	for {
		r := lexer.next()
		if _, ok := identifierTrailing[r]; !ok {
			lexer.back()
			break
		}
	}
	value := lexer.expression[start:lexer.currentPos]
	return token{
		tokenType: tUnquotedIdentifier,
		value:     value,
		position:  start,
		length:    lexer.currentPos - start,
	}
}

func (lexer *Lexer) consumeNumber() token {
	// Consume runes until we reach something that's not a number.
	start := lexer.currentPos - lexer.lastWidth
	for {
		r := lexer.next()
		if r < '0' || r > '9' {
			lexer.back()
			break
		}
	}
	value := lexer.expression[start:lexer.currentPos]
	return token{
		tokenType: tNumber,
		value:     value,
		position:  start,
		length:    lexer.currentPos - start,
	}
}
<commit_msg>perf: Switch from map to bitmask for lexing identifiers<commit_after>package jmespath

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"unicode/utf8"
)

type token struct {
	tokenType tokType
	value     string
	position  int
	length    int
}

type tokType int

const eof = -1

// Lexer contains information about the expression being tokenized.
type Lexer struct {
	expression string       // The expression provided by the user.
	currentPos int          // The current position in the string.
	lastWidth  int          // The width of the current rune.  This
	buf        bytes.Buffer // Internal buffer used for building up values.
}

// SyntaxError is the main error used whenever a lexing or parsing error occurs.
type SyntaxError struct {
	msg        string // Error message displayed to user
	Expression string // Expression that generated a SyntaxError
	Offset     int    // The location in the string where the error occurred
}

func (e SyntaxError) Error() string {
	// In the future, it would be good to underline the specific
	// location where the error occurred.
	return e.msg
}

//go:generate stringer -type=tokType
const (
	tUnknown tokType = iota
	tStar
	tDot
	tFilter
	tFlatten
	tLparen
	tRparen
	tLbracket
	tRbracket
	tLbrace
	tRbrace
	tOr
	tPipe
	tNumber
	tUnquotedIdentifier
	tQuotedIdentifier
	tComma
	tColon
	tLT
	tLTE
	tGT
	tGTE
	tEQ
	tNE
	tJSONLiteral
	tStringLiteral
	tCurrent
	tExpref
	tEOF
)

var basicTokens = map[rune]tokType{
	'.': tDot,
	'*': tStar,
	',': tComma,
	':': tColon,
	'{': tLbrace,
	'}': tRbrace,
	']': tRbracket, // tLbracket not included because it could be "[]"
	'(': tLparen,
	')': tRparen,
	'@': tCurrent,
	'&': tExpref,
}

// Bit mask for [a-zA-Z_] shifted down 64 bits to fit in a single uint64.
// When using this bitmask just be sure to shift the rune down 64 bits
// before checking against identifierStartBits.
const identifierStartBits uint64 = 576460745995190270

// Bit mask for [a-zA-Z0-9], 128 bits -> 2 uint64s.
var identifierTrailingBits [2]uint64 = [2]uint64{287948901175001088, 576460745995190270}

var whiteSpace = map[rune]bool{
	' ': true, '\t': true, '\n': true, '\r': true,
}

func (t token) String() string {
	return fmt.Sprintf("Token{%+v, %s, %d, %d}",
		t.tokenType, t.value, t.position, t.length)
}

// NewLexer creates a new JMESPath lexer.
func NewLexer() *Lexer {
	lexer := Lexer{}
	return &lexer
}

func (lexer *Lexer) next() rune {
	if lexer.currentPos >= len(lexer.expression) {
		lexer.lastWidth = 0
		return eof
	}
	r, w := utf8.DecodeRuneInString(lexer.expression[lexer.currentPos:])
	lexer.lastWidth = w
	lexer.currentPos += w
	return r
}

func (lexer *Lexer) back() {
	lexer.currentPos -= lexer.lastWidth
}

func (lexer *Lexer) peek() rune {
	t := lexer.next()
	lexer.back()
	return t
}

// tokenize takes an expression and returns corresponding tokens.
func (lexer *Lexer) tokenize(expression string) ([]token, error) {
	var tokens []token
	lexer.expression = expression
	lexer.currentPos = 0
	lexer.lastWidth = 0
loop:
	for {
		r := lexer.next()
		if identifierStartBits&(1<<(uint64(r)-64)) > 0 {
			t := lexer.consumeUnquotedIdentifier()
			tokens = append(tokens, t)
		} else if val, ok := basicTokens[r]; ok {
			// Basic single char token.
			t := token{
				tokenType: val,
				value:     string(r),
				position:  lexer.currentPos - lexer.lastWidth,
				length:    1,
			}
			tokens = append(tokens, t)
		} else if r == '-' || (r >= '0' && r <= '9') {
			t := lexer.consumeNumber()
			tokens = append(tokens, t)
		} else if r == '[' {
			t := lexer.consumeLBracket()
			tokens = append(tokens, t)
		} else if r == '"' {
			t, err := lexer.consumeQuotedIdentifier()
			if err != nil {
				return tokens, err
			}
			tokens = append(tokens, t)
		} else if r == '\'' {
			t, err := lexer.consumeRawStringLiteral()
			if err != nil {
				return tokens, err
			}
			tokens = append(tokens, t)
		} else if r == '`' {
			t, err := lexer.consumeLiteral()
			if err != nil {
				return tokens, err
			}
			tokens = append(tokens, t)
		} else if r == '|' {
			t := lexer.matchOrElse(r, '|', tOr, tPipe)
			tokens = append(tokens, t)
		} else if r == '<' {
			t := lexer.matchOrElse(r, '=', tLTE, tLT)
			tokens = append(tokens, t)
		} else if r == '>' {
			t := lexer.matchOrElse(r, '=', tGTE, tGT)
			tokens = append(tokens, t)
		} else if r == '!' {
			t := lexer.matchOrElse(r, '=', tNE, tUnknown)
			tokens = append(tokens, t)
		} else if r == '=' {
			t := lexer.matchOrElse(r, '=', tEQ, tUnknown)
			tokens = append(tokens, t)
		} else if r == eof {
			break loop
		} else if _, ok := whiteSpace[r]; ok {
			// Ignore whitespace
		} else {
			return tokens, lexer.syntaxError(fmt.Sprintf("Unknown char: %s", strconv.QuoteRuneToASCII(r)))
		}
	}
	tokens = append(tokens, token{tEOF, "", len(lexer.expression), 0})
	return tokens, nil
}

// Consume characters until the ending rune "r" is reached.
// If the end of the expression is reached before seeing the
// terminating rune "r", then an error is returned.
// If no error occurs then the matching substring is returned.
// The returned string will not include the ending rune.
func (lexer *Lexer) consumeUntil(end rune) (string, error) {
	start := lexer.currentPos
	current := lexer.next()
	for current != end && current != eof {
		if current == '\\' && lexer.peek() != eof {
			lexer.next()
		}
		current = lexer.next()
	}
	if lexer.lastWidth == 0 {
		// Then we hit an EOF so we never reached the closing
		// delimiter.
		return "", &SyntaxError{
			msg:        "Unclosed delimiter: " + string(end),
			Expression: lexer.expression,
			Offset:     len(lexer.expression),
		}
	}
	return lexer.expression[start : lexer.currentPos-lexer.lastWidth], nil
}

func (lexer *Lexer) consumeLiteral() (token, error) {
	start := lexer.currentPos
	value, err := lexer.consumeUntil('`')
	if err != nil {
		return token{}, err
	}
	value = strings.Replace(value, "\\`", "`", -1)
	return token{
		tokenType: tJSONLiteral,
		value:     value,
		position:  start,
		length:    len(value),
	}, nil
}

func (lexer *Lexer) consumeRawStringLiteral() (token, error) {
	start := lexer.currentPos
	currentIndex := start
	current := lexer.next()
	for current != '\'' && lexer.peek() != eof {
		if current == '\\' && lexer.peek() == '\'' {
			chunk := lexer.expression[currentIndex : lexer.currentPos-1]
			lexer.buf.WriteString(chunk)
			lexer.buf.WriteString("'")
			lexer.next()
			currentIndex = lexer.currentPos
		}
		current = lexer.next()
	}
	if lexer.lastWidth == 0 {
		// Then we hit an EOF so we never reached the closing
		// delimiter.
		return token{}, &SyntaxError{
			msg:        "Unclosed delimiter: '",
			Expression: lexer.expression,
			Offset:     len(lexer.expression),
		}
	}
	if currentIndex < lexer.currentPos {
		lexer.buf.WriteString(lexer.expression[currentIndex : lexer.currentPos-1])
	}
	value := lexer.buf.String()
	// Reset the buffer so it can reused again.
	lexer.buf.Reset()
	return token{
		tokenType: tStringLiteral,
		value:     value,
		position:  start,
		length:    len(value),
	}, nil
}

func (lexer *Lexer) syntaxError(msg string) SyntaxError {
	return SyntaxError{
		msg:        msg,
		Expression: lexer.expression,
		Offset:     lexer.currentPos,
	}
}

// Checks for a two char token, otherwise matches a single character
// token. This is used whenever a two char token overlaps a single
// char token, e.g. "||" -> tPipe, "|" -> tOr.
func (lexer *Lexer) matchOrElse(first rune, second rune, matchedType tokType, singleCharType tokType) token {
	start := lexer.currentPos - lexer.lastWidth
	nextRune := lexer.next()
	var t token
	if nextRune == second {
		t = token{
			tokenType: matchedType,
			value:     string(first) + string(second),
			position:  start,
			length:    2,
		}
	} else {
		lexer.back()
		t = token{
			tokenType: singleCharType,
			value:     string(first),
			position:  start,
			length:    1,
		}
	}
	return t
}

func (lexer *Lexer) consumeLBracket() token {
	// There's three options here:
	// 1. A filter expression "[?"
	// 2. A flatten operator "[]"
	// 3. A bare rbracket "["
	start := lexer.currentPos - lexer.lastWidth
	nextRune := lexer.next()
	var t token
	if nextRune == '?' {
		t = token{
			tokenType: tFilter,
			value:     "[?",
			position:  start,
			length:    2,
		}
	} else if nextRune == ']' {
		t = token{
			tokenType: tFlatten,
			value:     "[]",
			position:  start,
			length:    2,
		}
	} else {
		t = token{
			tokenType: tLbracket,
			value:     "[",
			position:  start,
			length:    1,
		}
		lexer.back()
	}
	return t
}

func (lexer *Lexer) consumeQuotedIdentifier() (token, error) {
	start := lexer.currentPos
	value, err := lexer.consumeUntil('"')
	if err != nil {
		return token{}, err
	}
	var decoded string
	asJSON := []byte("\"" + value + "\"")
	if err := json.Unmarshal([]byte(asJSON), &decoded); err != nil {
		return token{}, err
	}
	return token{
		tokenType: tQuotedIdentifier,
		value:     decoded,
		position:  start - 1,
		length:    len(decoded),
	}, nil
}

func (lexer *Lexer) consumeUnquotedIdentifier() token {
	// Consume runes until we reach the end of an unquoted
	// identifier.
	start := lexer.currentPos - lexer.lastWidth
	for {
		r := lexer.next()
		if r < 0 || r > 128 || identifierTrailingBits[uint64(r)/64]&(1<<(uint64(r)%64)) == 0 {
			lexer.back()
			break
		}
	}
	value := lexer.expression[start:lexer.currentPos]
	return token{
		tokenType: tUnquotedIdentifier,
		value:     value,
		position:  start,
		length:    lexer.currentPos - start,
	}
}

func (lexer *Lexer) consumeNumber() token {
	// Consume runes until we reach something that's not a number.
	start := lexer.currentPos - lexer.lastWidth
	for {
		r := lexer.next()
		if r < '0' || r > '9' {
			lexer.back()
			break
		}
	}
	value := lexer.expression[start:lexer.currentPos]
	return token{
		tokenType: tNumber,
		value:     value,
		position:  start,
		length:    lexer.currentPos - start,
	}
}
<|endoftext|>


<commit_before>package cask

// The Lexer has been stolen from:
// https://github.com/goruby/goruby/blob/master/lexer

import (
	"bytes"
	"fmt"
	"unicode"
	"unicode/utf8"
)

const eof = -1

// Lexer is the engine to process input and emit Tokens.
type Lexer struct {
	// input specifies the string being scanned.
	input string

	// state specifies the next lexing function to enter.
	state StateFn

	// position specifies current position in the input.
	position int

	// lines specifies the number of lines that have been lexed in input.
	lines int

	// start specifies the position of this item.
	start int

	// width specifies the width of the last rune read from input.
	width int

	// tokens specifies the channel of scanned tokens.
	tokens chan Token
}

// LexStartFn represents the entrypoint the Lexer uses to start processing the
// input.
var LexStartFn = startLexer

// StateFn represents a function which is capable of lexing parts of the
// input. It returns another StateFn to proceed with.
//
// Typically a state function would get called from LexStartFn and should
// return LexStartFn to go back to the decision loop. It also could return
// another non start state function if the partial input to parse is abiguous.
type StateFn func(*Lexer) StateFn

// NewLexer creates a new Lexer instance and returns its pointer. Requires an
// input to be passed as an argument that is ready to be processed.
func NewLexer(input string) *Lexer {
	return &Lexer{
		input:  input,
		state:  startLexer,
		tokens: make(chan Token, 3), // three tokens are sufficient.
	}
}

// NextToken will return the next token processed from the lexer.
func (l *Lexer) NextToken() Token {
	for {
		select {
		case item, ok := <-l.tokens:
			if ok {
				return item
			}
			panic(fmt.Errorf("No tokens left"))
		default:
			l.state = l.state(l)
			if l.state == nil {
				close(l.tokens)
			}
		}
	}
}

// HasNext returns true if there are tokens left, false if EOF has reached.
func (l *Lexer) HasNext() bool {
	return l.state != nil
}

// emit passes a token back to the client.
func (l *Lexer) emit(t TokenType) {
	l.tokens <- *NewToken(t, l.input[l.start:l.position], l.start)
	l.start = l.position
}

// next returns the next rune in the input.
func (l *Lexer) next() rune {
	if l.position >= len(l.input) {
		l.width = 0
		return eof
	}
	var r rune
	r, l.width = utf8.DecodeRuneInString(l.input[l.position:])
	l.position += l.width
	return r
}

// ignore skips over the pending input before this point.
func (l *Lexer) ignore() {
	l.start = l.position
}

// backup steps back one rune.
func (l *Lexer) backup() {
	l.position -= l.width
}

// peek returns but does not consume the next rune in the input.
func (l *Lexer) peek() rune {
	r := l.next()
	l.backup()
	return r
}

// errorf returns an error token and terminates the scan by passing
// back a nil pointer that will be the next state, terminating l.run.
func (l *Lexer) errorf(format string, args ...interface{}) StateFn {
	l.tokens <- *NewToken(ILLEGAL, fmt.Sprintf(format, args...), l.start)
	return nil
}

// startLexer starts the Lexer processing.
func startLexer(l *Lexer) StateFn {
	r := l.next()
	if isWhitespace(r) {
		l.ignore()
		return startLexer
	}
	switch r {
	case '$':
		return lexGlobal
	case '\n':
		l.lines++
		l.emit(NEWLINE)
		return startLexer
	case '\'':
		return lexSingleQuoteString
	case '"':
		return lexString
	case ':':
		if p := l.peek(); p == ':' {
			l.next()
			l.emit(SCOPE)
			return startLexer
		}

		return lexSymbol
	case '.':
		l.emit(DOT)
		return startLexer
	case '=':
		if l.peek() == '=' {
			l.next()
			l.emit(EQ)
		} else {
			l.emit(ASSIGN)
		}

		return startLexer
	case '+':
		l.emit(PLUS)
		return startLexer
	case '-':
		l.emit(MINUS)
		return startLexer
	case '!':
		if l.peek() == '=' {
			l.next()
			l.emit(NOTEQ)
		} else {
			l.emit(BANG)
		}

		return startLexer
	case '/':
		l.emit(SLASH)
		return startLexer
	case '%':
		if l.peek() == 'r' {
			l.next()
			l.emit(PNREGEXP)

			return lexRegexp
		}

		l.emit(MODULUS)

		return startLexer
	case '*':
		l.emit(ASTERISK)
		return startLexer
	case '<':
		l.emit(LT)
		return startLexer
	case '>':
		l.emit(GT)
		return startLexer
	case '(':
		l.emit(LPAREN)
		return startLexer
	case ')':
		l.emit(RPAREN)
		return startLexer
	case '{':
		l.emit(LBRACE)
		return startLexer
	case '}':
		l.emit(RBRACE)
		return startLexer
	case '[':
		l.emit(LBRACKET)
		return startLexer
	case ']':
		l.emit(RBRACKET)
		return startLexer
	case ',':
		l.emit(COMMA)
		return startLexer
	case ';':
		l.emit(SEMICOLON)
		return startLexer
	case eof:
		l.emit(EOF)
		return startLexer
	case '#':
		return commentLexer
	case '|':
		l.emit(PIPE)
		return startLexer
	default:
		if isLetter(r) {
			return lexIdentifier
		}

		if isDigit(r) {
			return lexDigit
		}

		return l.errorf("Illegal character at %d: '%c'", l.start, r)
	}
}

// lexIdentifier lexes the identifier.
func lexIdentifier(l *Lexer) StateFn {
	legalIdentifierCharacters := []byte{'?', '!'}

	r := l.next()
	for isLetter(r) || isDigit(r) || bytes.ContainsRune(legalIdentifierCharacters, r) {
		r = l.next()
	}

	l.backup()
	literal := l.input[l.start:l.position]
	l.emit(LookupIdent(literal))

	return startLexer
}

// lexIdentifier lexes the digit.
func lexDigit(l *Lexer) StateFn {
	r := l.next()
	for isDigitOrUnderscore(r) {
		r = l.next()
	}

	l.backup()
	l.emit(INT)

	return startLexer
}

// lexSingleQuoteString lexes the single quote string.
func lexSingleQuoteString(l *Lexer) StateFn {
	l.ignore()
	r := l.next()

	for r != '\'' {
		r = l.next()
	}
	l.backup()
	l.emit(STRING)
	l.next()
	l.ignore()

	return startLexer
}

// lexString lexes the string.
func lexString(l *Lexer) StateFn {
	l.ignore()

	r := l.next()
	for r != '"' {
		r = l.next()
	}

	l.backup()
	l.emit(STRING)
	l.next()
	l.ignore()

	return startLexer
}

// lexRegexp lexes the regular expression.
func lexRegexp(l *Lexer) StateFn {
	l.ignore()

	r := l.next()
	l.emit(PNSTART)

	pnStart := r
	pnEnd := pnStart

	switch pnStart {
	case '(':
		pnEnd = ')'
		break
	case '[':
		pnEnd = ']'
		break
	case '{':
		pnEnd = '}'
		break
	case '<':
		pnEnd = '>'
		break
	}

	for r != pnEnd {
		r = l.next()
	}

	l.backup()
	l.emit(REGEXP)
	l.next()
	l.emit(PNEND)

	return startLexer
}

// lexSymbol lexes the symbol.
func lexSymbol(l *Lexer) StateFn {
	l.ignore()

	r := l.next()
	for isLetter(r) || isDigit(r) {
		r = l.next()
	}

	l.backup()
	l.emit(SYMBOL)

	return startLexer
}

// lexGlobal lexes the global.
func lexGlobal(l *Lexer) StateFn {
	r := l.next()

	if isExpressionDelimiter(r) {
		return l.errorf("Illegal character at %d: '%c'", l.position, r)
	}

	if isWhitespace(r) {
		return l.errorf("Illegal character at %d: '%c'", l.position, r)
	}

	for !isWhitespace(r) && !isExpressionDelimiter(r) {
		r = l.next()
	}

	l.backup()
	l.emit(GLOBAL)

	return startLexer
}

// commentLexer lexes the comment.
func commentLexer(l *Lexer) StateFn {
	r := l.next()
	for r != '\n' {
		r = l.next()
	}

	l.ignore()

	return startLexer
}

// isDigit checks whether the specified rune is a whitespace character.
func isWhitespace(r rune) bool {
	return unicode.IsSpace(r) && r != '\n'
}

// isDigit checks whether the specified rune is a letter.
func isLetter(r rune) bool {
	return 'a' <= r && r <= 'z' || 'A' <= r && r <= 'Z' || r == '_'
}

// isDigit checks whether the specified rune is a digit.
func isDigit(r rune) bool {
	return '0' <= r && r <= '9'
}

// isDigitOrUnderscore checks whether the specified rune is a digit or an
// underscore.
func isDigitOrUnderscore(r rune) bool {
	return isDigit(r) || r == '_'
}

// isExpressionDelimiter checks whether the specified rune is an expression
// delimiter.
func isExpressionDelimiter(r rune) bool {
	return r == '\n' || r == ';' || r == eof
}
<commit_msg>Add escape characters support in lexString() and lexSingleQuoteString()<commit_after>package cask

// The Lexer has been stolen from:
// https://github.com/goruby/goruby/blob/master/lexer

import (
	"bytes"
	"fmt"
	"unicode"
	"unicode/utf8"
)

const eof = -1

// Lexer is the engine to process input and emit Tokens.
type Lexer struct {
	// input specifies the string being scanned.
	input string

	// state specifies the next lexing function to enter.
	state StateFn

	// position specifies current position in the input.
	position int

	// lines specifies the number of lines that have been lexed in input.
	lines int

	// start specifies the position of this item.
	start int

	// width specifies the width of the last rune read from input.
	width int

	// tokens specifies the channel of scanned tokens.
	tokens chan Token
}

// LexStartFn represents the entrypoint the Lexer uses to start processing the
// input.
var LexStartFn = startLexer

// StateFn represents a function which is capable of lexing parts of the
// input. It returns another StateFn to proceed with.
//
// Typically a state function would get called from LexStartFn and should
// return LexStartFn to go back to the decision loop. It also could return
// another non start state function if the partial input to parse is abiguous.
type StateFn func(*Lexer) StateFn

// NewLexer creates a new Lexer instance and returns its pointer. Requires an
// input to be passed as an argument that is ready to be processed.
func NewLexer(input string) *Lexer {
	return &Lexer{
		input:  input,
		state:  startLexer,
		tokens: make(chan Token, 3), // three tokens are sufficient.
	}
}

// NextToken will return the next token processed from the lexer.
func (l *Lexer) NextToken() Token {
	for {
		select {
		case item, ok := <-l.tokens:
			if ok {
				return item
			}
			panic(fmt.Errorf("No tokens left"))
		default:
			l.state = l.state(l)
			if l.state == nil {
				close(l.tokens)
			}
		}
	}
}

// HasNext returns true if there are tokens left, false if EOF has reached.
func (l *Lexer) HasNext() bool {
	return l.state != nil
}

// emit passes a token back to the client.
func (l *Lexer) emit(t TokenType) {
	l.tokens <- *NewToken(t, l.input[l.start:l.position], l.start)
	l.start = l.position
}

// next returns the next rune in the input.
func (l *Lexer) next() rune {
	if l.position >= len(l.input) {
		l.width = 0
		return eof
	}

	var r rune
	r, l.width = utf8.DecodeRuneInString(l.input[l.position:])
	l.position += l.width
	return r
}

// ignore skips over the pending input before this point.
func (l *Lexer) ignore() {
	l.start = l.position
}

// backup steps back one rune.
func (l *Lexer) backup() {
	l.position -= l.width
}

// peek returns but does not consume the next rune in the input.
func (l *Lexer) peek() rune {
	r := l.next()
	l.backup()
	return r
}

// errorf returns an error token and terminates the scan by passing
// back a nil pointer that will be the next state, terminating l.run.
func (l *Lexer) errorf(format string, args ...interface{}) StateFn {
	l.tokens <- *NewToken(ILLEGAL, fmt.Sprintf(format, args...), l.start)
	return nil
}

// startLexer starts the Lexer processing.
func startLexer(l *Lexer) StateFn {
	r := l.next()
	if isWhitespace(r) {
		l.ignore()
		return startLexer
	}
	switch r {
	case '$':
		return lexGlobal
	case '\n':
		l.lines++
		l.emit(NEWLINE)
		return startLexer
	case '\'':
		return lexSingleQuoteString
	case '"':
		return lexString
	case ':':
		if p := l.peek(); p == ':' {
			l.next()
			l.emit(SCOPE)
			return startLexer
		}

		return lexSymbol
	case '.':
		l.emit(DOT)
		return startLexer
	case '=':
		if l.peek() == '=' {
			l.next()
			l.emit(EQ)
		} else {
			l.emit(ASSIGN)
		}

		return startLexer
	case '+':
		l.emit(PLUS)
		return startLexer
	case '-':
		l.emit(MINUS)
		return startLexer
	case '!':
		if l.peek() == '=' {
			l.next()
			l.emit(NOTEQ)
		} else {
			l.emit(BANG)
		}

		return startLexer
	case '/':
		l.emit(SLASH)
		return startLexer
	case '%':
		if l.peek() == 'r' {
			l.next()
			l.emit(PNREGEXP)

			return lexRegexp
		}

		l.emit(MODULUS)

		return startLexer
	case '*':
		l.emit(ASTERISK)
		return startLexer
	case '<':
		l.emit(LT)
		return startLexer
	case '>':
		l.emit(GT)
		return startLexer
	case '(':
		l.emit(LPAREN)
		return startLexer
	case ')':
		l.emit(RPAREN)
		return startLexer
	case '{':
		l.emit(LBRACE)
		return startLexer
	case '}':
		l.emit(RBRACE)
		return startLexer
	case '[':
		l.emit(LBRACKET)
		return startLexer
	case ']':
		l.emit(RBRACKET)
		return startLexer
	case ',':
		l.emit(COMMA)
		return startLexer
	case ';':
		l.emit(SEMICOLON)
		return startLexer
	case eof:
		l.emit(EOF)
		return startLexer
	case '#':
		return commentLexer
	case '|':
		l.emit(PIPE)
		return startLexer
	default:
		if isLetter(r) {
			return lexIdentifier
		}

		if isDigit(r) {
			return lexDigit
		}

		return l.errorf("Illegal character at %d: '%c'", l.start, r)
	}
}

// lexIdentifier lexes the identifier.
func lexIdentifier(l *Lexer) StateFn {
	legalIdentifierCharacters := []byte{'?', '!'}

	r := l.next()
	for isLetter(r) || isDigit(r) || bytes.ContainsRune(legalIdentifierCharacters, r) {
		r = l.next()
	}

	l.backup()
	literal := l.input[l.start:l.position]
	l.emit(LookupIdent(literal))

	return startLexer
}

// lexIdentifier lexes the digit.
func lexDigit(l *Lexer) StateFn {
	r := l.next()
	for isDigitOrUnderscore(r) {
		r = l.next()
	}

	l.backup()
	l.emit(INT)

	return startLexer
}

// lexSingleQuoteString lexes the single quote string.
func lexSingleQuoteString(l *Lexer) StateFn {
	l.ignore()
	r := l.next()

	for r != '\'' {
		// we skip the escape character
		if l.peek() == '\'' && r == '\\' {
			l.next()
		}

		r = l.next()
	}

	l.backup()
	l.emit(STRING)
	l.next()
	l.ignore()

	return startLexer
}

// lexString lexes the string.
func lexString(l *Lexer) StateFn {
	l.ignore()
	r := l.next()

	for r != '"' {
		// we skip the escape character
		if l.peek() == '"' && r == '\\' {
			l.next()
		}

		r = l.next()
	}

	l.backup()
	l.emit(STRING)
	l.next()
	l.ignore()

	return startLexer
}

// lexRegexp lexes the regular expression.
func lexRegexp(l *Lexer) StateFn {
	l.ignore()

	r := l.next()
	l.emit(PNSTART)

	pnStart := r
	pnEnd := pnStart

	switch pnStart {
	case '(':
		pnEnd = ')'
		break
	case '[':
		pnEnd = ']'
		break
	case '{':
		pnEnd = '}'
		break
	case '<':
		pnEnd = '>'
		break
	}

	for r != pnEnd {
		r = l.next()
	}

	l.backup()
	l.emit(REGEXP)
	l.next()
	l.emit(PNEND)

	return startLexer
}

// lexSymbol lexes the symbol.
func lexSymbol(l *Lexer) StateFn {
	l.ignore()

	r := l.next()
	for isLetter(r) || isDigit(r) {
		r = l.next()
	}

	l.backup()
	l.emit(SYMBOL)

	return startLexer
}

// lexGlobal lexes the global.
func lexGlobal(l *Lexer) StateFn {
	r := l.next()

	if isExpressionDelimiter(r) {
		return l.errorf("Illegal character at %d: '%c'", l.position, r)
	}

	if isWhitespace(r) {
		return l.errorf("Illegal character at %d: '%c'", l.position, r)
	}

	for !isWhitespace(r) && !isExpressionDelimiter(r) {
		r = l.next()
	}

	l.backup()
	l.emit(GLOBAL)

	return startLexer
}

// commentLexer lexes the comment.
func commentLexer(l *Lexer) StateFn {
	r := l.next()
	for r != '\n' {
		r = l.next()
	}

	l.ignore()

	return startLexer
}

// isDigit checks whether the specified rune is a whitespace character.
func isWhitespace(r rune) bool {
	return unicode.IsSpace(r) && r != '\n'
}

// isDigit checks whether the specified rune is a letter.
func isLetter(r rune) bool {
	return 'a' <= r && r <= 'z' || 'A' <= r && r <= 'Z' || r == '_'
}

// isDigit checks whether the specified rune is a digit.
func isDigit(r rune) bool {
	return '0' <= r && r <= '9'
}

// isDigitOrUnderscore checks whether the specified rune is a digit or an
// underscore.
func isDigitOrUnderscore(r rune) bool {
	return isDigit(r) || r == '_'
}

// isExpressionDelimiter checks whether the specified rune is an expression
// delimiter.
func isExpressionDelimiter(r rune) bool {
	return r == '\n' || r == ';' || r == eof
}
<|endoftext|>


<commit_before>package easyss

import (
	"fmt"
	"io"
	"net"
	"strconv"

	"github.com/nange/easypool"
	"github.com/nange/easyss/cipherstream"
	"github.com/nange/easyss/util"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"github.com/txthinking/socks5"
)

var dataHeaderBytes = util.NewBytes(util.Http2HeaderLen)

func (ss *Easyss) LocalSocks5() error {
	var addr string
	if ss.BindAll() {
		addr = ":" + strconv.Itoa(ss.LocalPort())
	} else {
		addr = "127.0.0.1:" + strconv.Itoa(ss.LocalPort())
	}
	log.Infof("starting local socks5 server at %v", addr)

	//socks5.Debug = true
	server, err := socks5.NewClassicServer(addr, "127.0.0.1", "", "", 0, int(ss.Timeout()))
	if err != nil {
		log.Errorf("new socks5 server err: %+v", err)
		return err
	}
	ss.SetSocksServer(server)

	log.Warnf("local socks5 server:%s", server.ListenAndServe(ss))

	return nil
}

func (ss *Easyss) TCPHandle(s *socks5.Server, conn *net.TCPConn, r *socks5.Request) error {
	targetAddr := r.Address()
	log.Infof("target addr:%v", targetAddr)

	if err := ss.validateAddr(r.Address()); err != nil {
		log.Errorf("validate socks5 request:%v", err)
		return err
	}

	if r.Cmd == socks5.CmdConnect {
		a, addr, port, err := socks5.ParseAddress(ss.LocalAddr())
		if err != nil {
			log.Errorf("socks5 ParseAddress err:%+v", err)
			return err
		}
		p := socks5.NewReply(socks5.RepSuccess, a, addr, port)
		if _, err := p.WriteTo(conn); err != nil {
			return err
		}

		return ss.localRelay(conn, targetAddr)
	}

	if r.Cmd == socks5.CmdUDP {
		caddr, err := r.UDP(conn, s.UDPAddr)
		log.Infof("target request is udp proto, target addr:%v, caddr:%v, conn.LocalAddr:%s, conn.RemoteAddr:%s, s.UDPAddr:%v",
			targetAddr, caddr.String(), conn.LocalAddr().String(), conn.RemoteAddr().String(), s.ServerAddr)
		if err != nil {
			return err
		}

		// if client udp addr isn't private ip, we do not set associated udp
		// this case may be fired by non-standard socks5 implements
		if caddr.IP.IsLoopback() || caddr.IP.IsPrivate() || caddr.IP.IsUnspecified() {
			ch := make(chan byte)
			portStr := strconv.FormatInt(int64(caddr.Port), 10)
			s.AssociatedUDP.Set(portStr, ch, -1)
			defer func() {
				log.Debugf("exit associate tcp connection, close chan=========")
				close(ch)
				s.AssociatedUDP.Delete(portStr)
			}()
		}

		io.Copy(io.Discard, conn)
		log.Infof("A tcp connection that udp %v associated closed, target addr:%v\n", caddr.String(), targetAddr)
		return nil
	}

	return socks5.ErrUnsupportCmd
}

var paddingBytes = util.NewBytes(cipherstream.PaddingSize)

func (ss *Easyss) localRelay(localConn net.Conn, addr string) (err error) {
	pool := ss.Pool()
	if pool == nil {
		return errors.New("easyss is closed")
	}

	stream, err := pool.Get()
	if err != nil {
		log.Errorf("get stream from pool failed:%+v", err)
		return
	}

	log.Debugf("after pool get: current tcp pool has %v connections", pool.Len())
	defer func() {
		stream.Close()
		log.Debugf("after stream close: current tcp pool has %v connections", pool.Len())
	}()

	if err = ss.handShakeWithRemote(stream, addr, "tcp"); err != nil {
		log.Errorf("hand-shake with remote server err:%v", err)
		if pc, ok := stream.(*easypool.PoolConn); ok {
			log.Debugf("mark pool conn stream unusable")
			pc.MarkUnusable()
		}
		return
	}

	csStream, err := cipherstream.New(stream, ss.Password(), ss.Method(), "tcp")
	if err != nil {
		log.Errorf("new cipherstream err:%+v, password:%v, method:%v",
			err, ss.Password(), ss.Method())
		return
	}

	n1, n2, needClose := ss.relay(csStream, localConn)
	csStream.(*cipherstream.CipherStream).Release()

	log.Debugf("send %v bytes to %v, and recive %v bytes", n1, addr, n2)
	if !needClose {
		log.Debugf("underlying connection is health, so reuse it")
	}

	ss.stat.BytesSend.Add(n1)
	ss.stat.BytesRecive.Add(n2)

	return
}

func (ss *Easyss) validateAddr(addr string) error {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return fmt.Errorf("invalid target address:%v, err:%v", addr, err)
	}

	if !util.IsIP(host) {
		if host == ss.Server() {
			return fmt.Errorf("target host equals to server host, which may caused infinite-loop")
		}
		return nil
	}

	if util.IsPrivateIP(host) {
		return fmt.Errorf("target host:%v is private ip, which is invalid", host)
	}
	for _, ip := range ss.serverIPs {
		if host == ip {
			return fmt.Errorf("target host:%v equals server host ip, which may caused infinite-loop", host)
		}
	}

	return nil
}

func (ss *Easyss) handShakeWithRemote(stream net.Conn, addr, protoType string) error {
	header := dataHeaderBytes.Get(util.Http2HeaderLen)
	defer dataHeaderBytes.Put(header)

	header = util.EncodeHTTP2DataFrameHeader(protoType, len(addr)+1, header)
	gcm, err := cipherstream.NewAes256GCM([]byte(ss.Password()))
	if err != nil {
		return fmt.Errorf("cipherstream.NewAes256GCM err:%s", err.Error())
	}

	headerCipher, err := gcm.Encrypt(header)
	if err != nil {
		return fmt.Errorf("gcm.Encrypt err:%s", err.Error())
	}
	cipherMethod := EncodeCipherMethod(ss.Method())
	if cipherMethod == 0 {
		return fmt.Errorf("unsupported cipher method:%s", ss.Method())
	}
	payloadCipher, err := gcm.Encrypt(append([]byte(addr), cipherMethod))
	if err != nil {
		return fmt.Errorf("gcm.Encrypt err:%s", err.Error())
	}

	handshake := append(headerCipher, payloadCipher...)
	if header[4] == 0x8 { // has padding field
		padBytes := paddingBytes.Get(cipherstream.PaddingSize)
		defer paddingBytes.Put(padBytes)

		var padCipher []byte
		padCipher, err = gcm.Encrypt(padBytes)
		if err != nil {
			return fmt.Errorf("encrypt padding buf err:%s", err.Error())
		}
		handshake = append(handshake, padCipher...)
	}
	_, err = stream.Write(handshake)

	return err
}

func EncodeCipherMethod(m string) byte {
	switch m {
	case "aes-256-gcm":
		return 1
	case "chacha20-poly1305":
		return 2
	default:
		return 0
	}
}
<commit_msg>local: modify some log-level for socks5 handshake<commit_after>package easyss

import (
	"fmt"
	"io"
	"net"
	"strconv"

	"github.com/nange/easypool"
	"github.com/nange/easyss/cipherstream"
	"github.com/nange/easyss/util"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"github.com/txthinking/socks5"
)

var dataHeaderBytes = util.NewBytes(util.Http2HeaderLen)

func (ss *Easyss) LocalSocks5() error {
	var addr string
	if ss.BindAll() {
		addr = ":" + strconv.Itoa(ss.LocalPort())
	} else {
		addr = "127.0.0.1:" + strconv.Itoa(ss.LocalPort())
	}
	log.Infof("starting local socks5 server at %v", addr)

	//socks5.Debug = true
	server, err := socks5.NewClassicServer(addr, "127.0.0.1", "", "", 0, int(ss.Timeout()))
	if err != nil {
		log.Errorf("new socks5 server err: %+v", err)
		return err
	}
	ss.SetSocksServer(server)

	log.Warnf("local socks5 server:%s", server.ListenAndServe(ss))

	return nil
}

func (ss *Easyss) TCPHandle(s *socks5.Server, conn *net.TCPConn, r *socks5.Request) error {
	targetAddr := r.Address()
	log.Infof("target addr:%v, is udp:%v", targetAddr, r.Cmd == socks5.CmdUDP)

	if err := ss.validateAddr(r.Address()); err != nil {
		log.Errorf("validate socks5 request:%v", err)
		return err
	}

	if r.Cmd == socks5.CmdConnect {
		a, addr, port, err := socks5.ParseAddress(ss.LocalAddr())
		if err != nil {
			log.Errorf("socks5 ParseAddress err:%+v", err)
			return err
		}
		p := socks5.NewReply(socks5.RepSuccess, a, addr, port)
		if _, err := p.WriteTo(conn); err != nil {
			return err
		}

		return ss.localRelay(conn, targetAddr)
	}

	if r.Cmd == socks5.CmdUDP {
		caddr, err := r.UDP(conn, s.UDPAddr)
		log.Debugf("target request is udp proto, target addr:%v, caddr:%v, conn.LocalAddr:%s, conn.RemoteAddr:%s, s.UDPAddr:%v",
			targetAddr, caddr.String(), conn.LocalAddr().String(), conn.RemoteAddr().String(), s.ServerAddr)
		if err != nil {
			return err
		}

		// if client udp addr isn't private ip, we do not set associated udp
		// this case may be fired by non-standard socks5 implements
		if caddr.IP.IsLoopback() || caddr.IP.IsPrivate() || caddr.IP.IsUnspecified() {
			ch := make(chan byte)
			portStr := strconv.FormatInt(int64(caddr.Port), 10)
			s.AssociatedUDP.Set(portStr, ch, -1)
			defer func() {
				log.Debugf("exit associate tcp connection, close chan=========")
				close(ch)
				s.AssociatedUDP.Delete(portStr)
			}()
		}

		io.Copy(io.Discard, conn)
		log.Debugf("A tcp connection that udp %v associated closed, target addr:%v\n", caddr.String(), targetAddr)
		return nil
	}

	return socks5.ErrUnsupportCmd
}

var paddingBytes = util.NewBytes(cipherstream.PaddingSize)

func (ss *Easyss) localRelay(localConn net.Conn, addr string) (err error) {
	pool := ss.Pool()
	if pool == nil {
		return errors.New("easyss is closed")
	}

	stream, err := pool.Get()
	if err != nil {
		log.Errorf("get stream from pool failed:%+v", err)
		return
	}

	log.Debugf("after pool get: current tcp pool has %v connections", pool.Len())
	defer func() {
		stream.Close()
		log.Debugf("after stream close: current tcp pool has %v connections", pool.Len())
	}()

	if err = ss.handShakeWithRemote(stream, addr, "tcp"); err != nil {
		log.Errorf("hand-shake with remote server err:%v", err)
		if pc, ok := stream.(*easypool.PoolConn); ok {
			log.Debugf("mark pool conn stream unusable")
			pc.MarkUnusable()
		}
		return
	}

	csStream, err := cipherstream.New(stream, ss.Password(), ss.Method(), "tcp")
	if err != nil {
		log.Errorf("new cipherstream err:%+v, password:%v, method:%v",
			err, ss.Password(), ss.Method())
		return
	}

	n1, n2, needClose := ss.relay(csStream, localConn)
	csStream.(*cipherstream.CipherStream).Release()

	log.Debugf("send %v bytes to %v, and recive %v bytes", n1, addr, n2)
	if !needClose {
		log.Debugf("underlying connection is health, so reuse it")
	}

	ss.stat.BytesSend.Add(n1)
	ss.stat.BytesRecive.Add(n2)

	return
}

func (ss *Easyss) validateAddr(addr string) error {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return fmt.Errorf("invalid target address:%v, err:%v", addr, err)
	}

	if !util.IsIP(host) {
		if host == ss.Server() {
			return fmt.Errorf("target host equals to server host, which may caused infinite-loop")
		}
		return nil
	}

	if util.IsPrivateIP(host) {
		return fmt.Errorf("target host:%v is private ip, which is invalid", host)
	}
	for _, ip := range ss.serverIPs {
		if host == ip {
			return fmt.Errorf("target host:%v equals server host ip, which may caused infinite-loop", host)
		}
	}

	return nil
}

func (ss *Easyss) handShakeWithRemote(stream net.Conn, addr, protoType string) error {
	header := dataHeaderBytes.Get(util.Http2HeaderLen)
	defer dataHeaderBytes.Put(header)

	header = util.EncodeHTTP2DataFrameHeader(protoType, len(addr)+1, header)
	gcm, err := cipherstream.NewAes256GCM([]byte(ss.Password()))
	if err != nil {
		return fmt.Errorf("cipherstream.NewAes256GCM err:%s", err.Error())
	}

	headerCipher, err := gcm.Encrypt(header)
	if err != nil {
		return fmt.Errorf("gcm.Encrypt err:%s", err.Error())
	}
	cipherMethod := EncodeCipherMethod(ss.Method())
	if cipherMethod == 0 {
		return fmt.Errorf("unsupported cipher method:%s", ss.Method())
	}
	payloadCipher, err := gcm.Encrypt(append([]byte(addr), cipherMethod))
	if err != nil {
		return fmt.Errorf("gcm.Encrypt err:%s", err.Error())
	}

	handshake := append(headerCipher, payloadCipher...)
	if header[4] == 0x8 { // has padding field
		padBytes := paddingBytes.Get(cipherstream.PaddingSize)
		defer paddingBytes.Put(padBytes)

		var padCipher []byte
		padCipher, err = gcm.Encrypt(padBytes)
		if err != nil {
			return fmt.Errorf("encrypt padding buf err:%s", err.Error())
		}
		handshake = append(handshake, padCipher...)
	}
	_, err = stream.Write(handshake)

	return err
}

func EncodeCipherMethod(m string) byte {
	switch m {
	case "aes-256-gcm":
		return 1
	case "chacha20-poly1305":
		return 2
	default:
		return 0
	}
}
<|endoftext|>


<commit_before>package locus

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"strings"
)

const cityUrl string = `http://api.ipinfodb.com/v3/ip-city/`
const countryUrl string = `http://api.ipinfodb.com/v3/ip-country/`

type Location struct {
	StatusCode    string `json: "statusCode"`
	StatusMessage string `json: "statusMessage"`
	IpAddress     string `json: "ipAddress"`
	CountryCode   string `json: "countryCode"`
	CountryName   string `json: "countryName"`
	RegionName    string `json: "regionName"`
	CityName      string `json: "cityName"`
	ZipCode       string `json: "zipCode"`
	Latitude      string `json: "latitude"`
	Longitude     string `json: "longitude"`
	TimeZone      string `json: "timeZone"`
}

func locationJson(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return body, nil
}

func location(url string) (*Location, error) {
	location := &Location{}

	raw_json, err := locationJson(url)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(raw_json, location)
	if err != nil {
		return nil, err
	}

	return location, nil
}

func requestUrl(ip string, precision string, key string) string {
	baseUrl := countryUrl
	if strings.ToLower(precision) == "city" {
		baseUrl = cityUrl
	}

	return strings.Join([]string{baseUrl, `?format=json`, `&ip=`, ip, `&key=`, key}, ``)
}

// Public API

func LookupLocationJson(ip string, precision string, key string) (string, error) {
	location, err := locationJson(requestUrl(ip, precision, key))
	return string(location[:]), err
}

func LookupLocation(ip string, precision string, key string) (*Location, error) {
	return location(requestUrl(ip, precision, key))
}

func BulkLookupLocationJSON(ips []string, precision string, key string) ([]string, error) {
	locations := make([]string, len(ips))
	var err error
	for i, ip := range ips {
		locations[i], err = LookupLocationJson(ip, precision, key)
		if err != nil {
			return nil, err
		}
	}

	return locations, nil
}

func BulkLookupLocation(ips []string, precision string, key string) ([]*Location, error) {
	locations := make([]*Location, len(ips))
	var err error
	for i, ip := range ips {
		locations[i], err = LookupLocation(ip, precision, key)
		if err != nil {
			return nil, err
		}
	}

	return locations, nil
}
<commit_msg>Simplify interface and implementation<commit_after>package locus

import (
	"bufio"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"
)

const cityUrl string = `http://api.ipinfodb.com/v3/ip-city/`
const countryUrl string = `http://api.ipinfodb.com/v3/ip-country/`

type Location struct {
	StatusCode    string `json: "statusCode"`
	StatusMessage string `json: "statusMessage"`
	IpAddress     string `json: "ipAddress"`
	CountryCode   string `json: "countryCode"`
	CountryName   string `json: "countryName"`
	RegionName    string `json: "regionName"`
	CityName      string `json: "cityName"`
	ZipCode       string `json: "zipCode"`
	Latitude      string `json: "latitude"`
	Longitude     string `json: "longitude"`
	TimeZone      string `json: "timeZone"`
}

func requestUrl(ip string, precision string, key string) (string, error) {
	baseUrl := countryUrl
	if strings.ToLower(precision) == "city" {
		baseUrl = cityUrl
	}

	var request *url.URL
	request, err := url.Parse(baseUrl)
	if err != nil {
		return ``, err
	}

	params := url.Values{}
	params.Set(`ip`, ip)
	params.Set(`key`, key)
	params.Set(`format`, `json`)
	request.RawQuery = params.Encode()

	return request.String(), nil
}

func lookupLocation(ip string, precision string, key string) (Location, error) {
	location := Location{}
	request, err := requestUrl(ip, precision, key)
	if err != nil {
		return Location{}, err
	}

	resp, err := http.Get(request)
	if err != nil {
		return Location{}, err
	}
	defer resp.Body.Close()

	raw_json, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return Location{}, err
	}

	err = json.Unmarshal(raw_json, &location)
	if err != nil {
		return Location{}, err
	}

	return location, nil
}

func lookupLocations(ips []string, precision string, key string) ([]Location, error) {
	locations := make([]Location, len(ips))
	var err error
	for i, ip := range ips {
		locations[i], err = LookupLocation(ip, precision, key)
		if err != nil {
			return nil, err
		}
	}

	return locations, nil
}

func lookupLocationsFile(filename string, precision string, key string) ([]Location, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	ips := make([]string, 0)
	for scanner := bufio.NewScanner(file); scanner.Scan(); {
		if scanner.Err() != nil {
			return nil, scanner.Err()
		}

		ips = append(ips, scanner.Text())
	}

	return LookupLocations(ips, precision, key)
}

// Public API

func LookupLocation(ip string, precision string, key string) (Location, error) {
	return lookupLocation(ip, precision, key)
}

func LookupLocations(ips []string, precision string, key string) ([]Location, error) {
	return lookupLocations(ips, precision, key)
}

func LookupLocationsFile(filename string, precision string, key string) ([]Location, error) {
	return lookupLocationsFile(filename, precision, key)
}
<|endoftext|>


<commit_before>/*
Logos

This program analyzes articles, identifying important elements in highly ranking articles.
*/

import (
  "io"
  "strings"
  "unicode"
)

type Publication struct {
  Score float64
  Author string
  Text PublicationBody
}

type PublicationBody interface {
  NextLine() (string)
  HasNextLine() (bool)
  NextWord() (string)
  HasNextWord() (bool)
  io.Reader
  io.Seeker
}

/*
---------------------------------------------------------------
-------------- Measures of Publication Quality ----------------
---------------------------------------------------------------
*/

func (p PublicationBody) WordCount() (int) {
  count := 0

  for (p.HasNextWord()) {
    l := p.NextWord()
    count++
  }

  return count
}

func (p PublicationBody) AverageWordsPerLine() (int) {
  sum := 0
  count := 0

  for (p.HasNextLine()) {
    l := p.NextLine()
    words := splitWords(l)
    sum += len(words)
    count++
  }

  return float64(sum) / count
}

func (p PublicationBody) AverageWordLength() (int) {
  sum := 0
  count := 0

  for (p.HasNextWord()) {
    w := p.NextWord()
    sum += len(w)
    count++
  }

  return float64(sum) / count
}

func (p PublicationBody) WordsLongerThan(x int) (int) {
  count := 0

  for (p.HasNextWord()) {
    w := p.NextWord()
    if (len(w) > x) {
      count++
    }
  }

  return count
}

func (p PublicationBody) WordsIn(list WordList) (int) {
  count := 0

  for (p.HasNextWord()) {
    w := p.NextWord()
    if (list.Contains(w)) {
      count++
    }
  }

  return count
}

func (p PublicationBody) ConstructMarkovMatrix(ngramSize int) (MarkovMatrix) {
  prevWords := make([]string, ngramSize)

  for (p.HasNextWord()) {
    w := p.NextWord()
    prevWords = append(prevWords, w)
  }
}

type NGram struct {
  Size int
  Words []string
}

type MarkovMatrix struct {
  Matrix map[NGram](map[NGram]float64)
}

type WordList struct {
  Words map[string]bool
}

func ConstructWordList(words []string) (WordList) {
  list := make(map[string]bool)

  for _, w := range words {
    list[w] = true
  }

  return list
}

func (w WordList) Contains(word string) (bool) {
  return w.Words[word]
}

/*
splitWords returns the separate words that make up a particular string, making
sure to remove punctuation and spaces.
*/
func splitWords(line string) ([]string) {
  f := func(c rune) bool {
    return unicode.IsPunct(c) || unicode.IsSpace(c)
  }
  return strings.FieldsFunc(sentence, f)
}
<commit_msg>Working on markov matrix stuff.<commit_after>/*
Logos

This program analyzes articles, identifying important elements in highly ranking articles.
*/

import (
  "io"
  "strings"
  "unicode"
  "strconv"
)

type Publication struct {
  Score float64
  Author string
  Text PublicationBody
}

type PublicationBody interface {
  NextLine() (string)
  HasNextLine() (bool)
  NextWord() (string)
  HasNextWord() (bool)
  io.Reader
  io.Seeker
}

/*
---------------------------------------------------------------
-------------- Measures of Publication Quality ----------------
---------------------------------------------------------------
*/

func (p PublicationBody) WordCount() (int) {
  count := 0

  for (p.HasNextWord()) {
    l := p.NextWord()
    count++
  }

  return count
}

func (p PublicationBody) AverageWordsPerLine() (int) {
  sum := 0
  count := 0

  for (p.HasNextLine()) {
    l := p.NextLine()
    words := splitWords(l)
    sum += len(words)
    count++
  }

  return float64(sum) / count
}

func (p PublicationBody) AverageWordLength() (int) {
  sum := 0
  count := 0

  for (p.HasNextWord()) {
    w := p.NextWord()
    sum += len(w)
    count++
  }

  return float64(sum) / count
}

func (p PublicationBody) WordsLongerThan(x int) (int) {
  count := 0

  for (p.HasNextWord()) {
    w := p.NextWord()
    if (len(w) > x) {
      count++
    }
  }

  return count
}

func (p PublicationBody) WordsIn(list WordList) (int) {
  count := 0

  for (p.HasNextWord()) {
    w := p.NextWord()
    if (list.Contains(w)) {
      count++
    }
  }

  return count
}

func (p PublicationBody) ConstructMarkovMatrix(ngramSize int) (MarkovMatrix) {
  prevWords := make([]string, ngramSize)

  for (p.HasNextWord()) {
    w := p.NextWord()
    prevWords = append(prevWords, w)
  }
}

type NGram struct {
  Size int
  Words []string
}

func (n NGram) Hash() (string) {
  s := []string{strconv.Itoa(n.Size)}
  s := append(s, n.Words...)
  return strings.Join(s, ".")
}

type MarkovMatrix struct {
  Matrix map[string]map[string]float64
}

func (m MarkovMatrix) SetProbability(i, j NGram, prob float64) {
  entry := m.Matrx[i.Hash()]
  if (entry == nil) {
    entry = make(map[string]float64)
  }
  entry[j.Hash()] = prob
  m.Matrix[i.Hash()] = entry
}

type WordList struct {
  Words map[string]bool
}

func ConstructWordList(words []string) (WordList) {
  list := make(map[string]bool)

  for _, w := range words {
    list[w] = true
  }

  return list
}

func (w WordList) Contains(word string) (bool) {
  return w.Words[word]
}

/*
splitWords returns the separate words that make up a particular string, making
sure to remove punctuation and spaces.
*/
func splitWords(line string) ([]string) {
  f := func(c rune) bool {
    return unicode.IsPunct(c) || unicode.IsSpace(c)
  }
  return strings.FieldsFunc(sentence, f)
}
<|endoftext|>