<commit_before>package gost

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"strconv"
	"time"

	"github.com/ginuerzh/gosocks5"
	"github.com/go-log/log"
	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
)

// Due to in/out byte length is inconsistent of the shadowsocks.Conn.Write,
// we wrap around it to make io.Copy happy.
type shadowConn struct {
	conn net.Conn
}

func (c *shadowConn) Read(b []byte) (n int, err error) {
	return c.conn.Read(b)
}

func (c *shadowConn) Write(b []byte) (n int, err error) {
	n = len(b) // force byte length consistent
	_, err = c.conn.Write(b)
	return
}

func (c *shadowConn) Close() error {
	return c.conn.Close()
}

func (c *shadowConn) LocalAddr() net.Addr {
	return c.conn.LocalAddr()
}

func (c *shadowConn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

func (c *shadowConn) SetDeadline(t time.Time) error {
	return c.conn.SetDeadline(t)
}

func (c *shadowConn) SetReadDeadline(t time.Time) error {
	return c.conn.SetReadDeadline(t)
}

func (c *shadowConn) SetWriteDeadline(t time.Time) error {
	return c.conn.SetWriteDeadline(t)
}

type shadowConnector struct {
	Cipher *url.Userinfo
}

// ShadowConnector creates a Connector for shadowsocks proxy client.
// It accepts a cipher info for shadowsocks data encryption/decryption.
// The cipher must not be nil.
func ShadowConnector(cipher *url.Userinfo) Connector {
	return &shadowConnector{Cipher: cipher}
}

func (c *shadowConnector) Connect(conn net.Conn, addr string) (net.Conn, error) {
	rawaddr, err := ss.RawAddr(addr)
	if err != nil {
		return nil, err
	}

	var method, password string
	if c.Cipher != nil {
		method = c.Cipher.Username()
		password, _ = c.Cipher.Password()
	}

	cipher, err := ss.NewCipher(method, password)
	if err != nil {
		return nil, err
	}

	sc, err := ss.DialWithRawAddrConn(rawaddr, conn, cipher)
	if err != nil {
		return nil, err
	}
	return &shadowConn{conn: sc}, nil
}

type shadowHandler struct {
	options *HandlerOptions
}

// ShadowHandler creates a server Handler for shadowsocks proxy server.
func ShadowHandler(opts ...HandlerOption) Handler {
	h := &shadowHandler{
		options: &HandlerOptions{},
	}
	for _, opt := range opts {
		opt(h.options)
	}
	return h
}

func (h *shadowHandler) Handle(conn net.Conn) {
	defer conn.Close()

	var method, password string
	users := h.options.Users
	if len(users) > 0 {
		method = users[0].Username()
		password, _ = users[0].Password()
	}
	cipher, err := ss.NewCipher(method, password)
	if err != nil {
		log.Log("[ss]", err)
		return
	}
	conn = &shadowConn{conn: ss.NewConn(conn, cipher)}

	log.Logf("[ss] %s - %s", conn.RemoteAddr(), conn.LocalAddr())

	addr, err := h.getRequest(conn)
	if err != nil {
		log.Logf("[ss] %s - %s : %s", conn.RemoteAddr(), conn.LocalAddr(), err)
		return
	}
	log.Logf("[ss] %s -> %s", conn.RemoteAddr(), addr)

	if !Can("tcp", addr, h.options.Whitelist, h.options.Blacklist) {
		log.Logf("[ss] Unauthorized to tcp connect to %s", addr)
		return
	}

	cc, err := h.options.Chain.Dial(addr)
	if err != nil {
		log.Logf("[ss] %s -> %s : %s", conn.RemoteAddr(), addr, err)
		return
	}
	defer cc.Close()

	log.Logf("[ss] %s <-> %s", conn.RemoteAddr(), addr)
	transport(conn, cc)
	log.Logf("[ss] %s >-< %s", conn.RemoteAddr(), addr)
}

const (
	idType  = 0 // address type index
	idIP0   = 1 // ip address start index
	idDmLen = 1 // domain address length index
	idDm0   = 2 // domain address start index

	typeIPv4 = 1 // type is ipv4 address
	typeDm   = 3 // type is domain address
	typeIPv6 = 4 // type is ipv6 address

	lenIPv4     = net.IPv4len + 2 // ipv4 + 2port
	lenIPv6     = net.IPv6len + 2 // ipv6 + 2port
	lenDmBase   = 2               // 1addrLen + 2port, plus addrLen
	lenHmacSha1 = 10
)

// This function is copied from shadowsocks library with some modification.
func (h *shadowHandler) getRequest(conn net.Conn) (host string, err error) {
	// buf size should at least have the same size with the largest possible
	// request size (when addrType is 3, domain name has at most 256 bytes)
	// 1(addrType) + 1(lenByte) + 256(max length address) + 2(port)
	buf := make([]byte, smallBufferSize)

	// read till we get possible domain length field
	conn.SetReadDeadline(time.Now().Add(ReadTimeout))
	if _, err = io.ReadFull(conn, buf[:idType+1]); err != nil {
		return
	}
	// clear timer
	conn.SetReadDeadline(time.Time{})

	var reqStart, reqEnd int
	addrType := buf[idType]
	switch addrType & ss.AddrMask {
	case typeIPv4:
		reqStart, reqEnd = idIP0, idIP0+lenIPv4
	case typeIPv6:
		reqStart, reqEnd = idIP0, idIP0+lenIPv6
	case typeDm:
		if _, err = io.ReadFull(conn, buf[idType+1:idDmLen+1]); err != nil {
			return
		}
		reqStart, reqEnd = idDm0, int(idDm0+buf[idDmLen]+lenDmBase)
	default:
		err = fmt.Errorf("addr type %d not supported", addrType&ss.AddrMask)
		return
	}

	if _, err = io.ReadFull(conn, buf[reqStart:reqEnd]); err != nil {
		return
	}

	// Return string for typeIP is not most efficient, but browsers (Chrome,
	// Safari, Firefox) all seems using typeDm exclusively. So this is not a
	// big problem.
	switch addrType & ss.AddrMask {
	case typeIPv4:
		host = net.IP(buf[idIP0 : idIP0+net.IPv4len]).String()
	case typeIPv6:
		host = net.IP(buf[idIP0 : idIP0+net.IPv6len]).String()
	case typeDm:
		host = string(buf[idDm0 : idDm0+buf[idDmLen]])
	}
	// parse port
	port := binary.BigEndian.Uint16(buf[reqEnd-2 : reqEnd])
	host = net.JoinHostPort(host, strconv.Itoa(int(port)))
	return
}

type shadowUDPListener struct {
	ln       net.PacketConn
	conns    map[string]*udpServerConn
	connChan chan net.Conn
	errChan  chan error
	ttl      time.Duration
}

// ShadowUDPListener creates a Listener for shadowsocks UDP relay server.
func ShadowUDPListener(addr string, cipher *url.Userinfo, ttl time.Duration) (Listener, error) {
	laddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	ln, err := net.ListenUDP("udp", laddr)
	if err != nil {
		return nil, err
	}

	var method, password string
	if cipher != nil {
		method = cipher.Username()
		password, _ = cipher.Password()
	}
	cp, err := ss.NewCipher(method, password)
	if err != nil {
		ln.Close()
		return nil, err
	}
	l := &shadowUDPListener{
		ln:       ss.NewSecurePacketConn(ln, cp, false),
		conns:    make(map[string]*udpServerConn),
		connChan: make(chan net.Conn, 1024),
		errChan:  make(chan error, 1),
		ttl:      ttl,
	}
	go l.listenLoop()
	return l, nil
}

func (l *shadowUDPListener) listenLoop() {
	for {
		b := make([]byte, mediumBufferSize)
		n, raddr, err := l.ln.ReadFrom(b)
		if err != nil {
			log.Logf("[ssu] peer -> %s : %s", l.Addr(), err)
			l.ln.Close()
			l.errChan <- err
			close(l.errChan)
			return
		}
		if Debug {
			log.Logf("[ssu] %s >>> %s : length %d", raddr, l.Addr(), n)
		}

		conn, ok := l.conns[raddr.String()]
		if !ok || conn.Closed() {
			conn = newUDPServerConn(l.ln, raddr, l.ttl)
			l.conns[raddr.String()] = conn

			select {
			case l.connChan <- conn:
			default:
				conn.Close()
				log.Logf("[ssu] %s - %s: connection queue is full", raddr, l.Addr())
			}
		}

		select {
		case conn.rChan <- b[:n]: // we keep the addr info so that the handler can identify the destination.
		default:
			log.Logf("[ssu] %s -> %s : read queue is full", raddr, l.Addr())
		}
	}
}

func (l *shadowUDPListener) Accept() (conn net.Conn, err error) {
	var ok bool
	select {
	case conn = <-l.connChan:
	case err, ok = <-l.errChan:
		if !ok {
			err = errors.New("accpet on closed listener")
		}
	}
	return
}

func (l *shadowUDPListener) Addr() net.Addr {
	return l.ln.LocalAddr()
}

func (l *shadowUDPListener) Close() error {
	return l.ln.Close()
}

type shadowUDPdHandler struct {
	ttl     time.Duration
	options *HandlerOptions
}

// ShadowUDPdHandler creates a server Handler for shadowsocks UDP relay server.
func ShadowUDPdHandler(opts ...HandlerOption) Handler {
	h := &shadowUDPdHandler{
		options: &HandlerOptions{},
	}
	for _, opt := range opts {
		opt(h.options)
	}
	return h
}

func (h *shadowUDPdHandler) Handle(conn net.Conn) {
	defer conn.Close()

	var err error
	var cc net.PacketConn
	if h.options.Chain.IsEmpty() {
		cc, err = net.ListenUDP("udp", nil)
		if err != nil {
			log.Logf("[ssu] %s - : %s", conn.LocalAddr(), err)
			return
		}
	} else {
		var c net.Conn
		c, err = getSOCKS5UDPTunnel(h.options.Chain, nil)
		if err != nil {
			log.Logf("[ssu] %s - : %s", conn.LocalAddr(), err)
			return
		}
		cc = &udpTunnelConn{Conn: c}
	}
	defer cc.Close()

	log.Logf("[ssu] %s <-> %s", conn.RemoteAddr(), conn.LocalAddr())
	transportUDP(conn, cc)
	log.Logf("[ssu] %s >-< %s", conn.RemoteAddr(), conn.LocalAddr())
}

func transportUDP(sc net.Conn, cc net.PacketConn) error {
	errc := make(chan error, 1)
	go func() {
		for {
			b := make([]byte, mediumBufferSize)
			n, err := sc.Read(b[3:]) // add rsv and frag fields to make it the standard SOCKS5 UDP datagram
			if err != nil {
				// log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), err)
				errc <- err
				return
			}
			dgram, err := gosocks5.ReadUDPDatagram(bytes.NewReader(b[:n+3]))
			if err != nil {
				log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), err)
				errc <- err
				return
			}
			//if Debug {
			//	log.Logf("[ssu] %s >>> %s length: %d", sc.RemoteAddr(), dgram.Header.Addr.String(), len(dgram.Data))
			//}
			addr, err := net.ResolveUDPAddr("udp", dgram.Header.Addr.String())
			if err != nil {
				errc <- err
				return
			}
			if _, err := cc.WriteTo(dgram.Data, addr); err != nil {
				errc <- err
				return
			}
		}
	}()

	go func() {
		for {
			b := make([]byte, mediumBufferSize)
			n, addr, err := cc.ReadFrom(b)
			if err != nil {
				errc <- err
				return
			}
			//if Debug {
			//	log.Logf("[ssu] %s <<< %s length: %d", sc.RemoteAddr(), addr, n)
			//}
			dgram := gosocks5.NewUDPDatagram(gosocks5.NewUDPHeader(0, 0, toSocksAddr(addr)), b[:n])
			buf := bytes.Buffer{}
			dgram.Write(&buf)
			if buf.Len() < 10 {
				log.Logf("[ssu] %s <- %s : invalid udp datagram", sc.RemoteAddr(), addr)
				continue
			}
			if _, err := sc.Write(buf.Bytes()[3:]); err != nil {
				errc <- err
				return
			}
		}
	}()

	err := <-errc
	if err != nil && err == io.EOF {
		err = nil
	}
	return err
}
<commit_msg>fix #165 #179<commit_after>package gost

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"strconv"
	"time"

	"github.com/ginuerzh/gosocks5"
	"github.com/go-log/log"
	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
)

// Due to in/out byte length is inconsistent of the shadowsocks.Conn.Write,
// we wrap around it to make io.Copy happy.
type shadowConn struct {
	conn net.Conn
}

func (c *shadowConn) Read(b []byte) (n int, err error) {
	return c.conn.Read(b)
}

func (c *shadowConn) Write(b []byte) (n int, err error) {
	n = len(b) // force byte length consistent
	_, err = c.conn.Write(b)
	return
}

func (c *shadowConn) Close() error {
	return c.conn.Close()
}

func (c *shadowConn) LocalAddr() net.Addr {
	return c.conn.LocalAddr()
}

func (c *shadowConn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

func (c *shadowConn) SetDeadline(t time.Time) error {
	return c.conn.SetDeadline(t)
}

func (c *shadowConn) SetReadDeadline(t time.Time) error {
	return c.conn.SetReadDeadline(t)
}

func (c *shadowConn) SetWriteDeadline(t time.Time) error {
	return c.conn.SetWriteDeadline(t)
}

type shadowConnector struct {
	Cipher *url.Userinfo
}

// ShadowConnector creates a Connector for shadowsocks proxy client.
// It accepts a cipher info for shadowsocks data encryption/decryption.
// The cipher must not be nil.
func ShadowConnector(cipher *url.Userinfo) Connector {
	return &shadowConnector{Cipher: cipher}
}

func (c *shadowConnector) Connect(conn net.Conn, addr string) (net.Conn, error) {
	rawaddr, err := ss.RawAddr(addr)
	if err != nil {
		return nil, err
	}

	var method, password string
	if c.Cipher != nil {
		method = c.Cipher.Username()
		password, _ = c.Cipher.Password()
	}

	cipher, err := ss.NewCipher(method, password)
	if err != nil {
		return nil, err
	}

	sc, err := ss.DialWithRawAddrConn(rawaddr, conn, cipher)
	if err != nil {
		return nil, err
	}
	return &shadowConn{conn: sc}, nil
}

type shadowHandler struct {
	options *HandlerOptions
}

// ShadowHandler creates a server Handler for shadowsocks proxy server.
func ShadowHandler(opts ...HandlerOption) Handler {
	h := &shadowHandler{
		options: &HandlerOptions{},
	}
	for _, opt := range opts {
		opt(h.options)
	}
	return h
}

func (h *shadowHandler) Handle(conn net.Conn) {
	defer conn.Close()

	var method, password string
	users := h.options.Users
	if len(users) > 0 {
		method = users[0].Username()
		password, _ = users[0].Password()
	}
	cipher, err := ss.NewCipher(method, password)
	if err != nil {
		log.Log("[ss]", err)
		return
	}
	conn = &shadowConn{conn: ss.NewConn(conn, cipher)}

	log.Logf("[ss] %s - %s", conn.RemoteAddr(), conn.LocalAddr())

	conn.SetReadDeadline(time.Now().Add(ReadTimeout))
	addr, err := h.getRequest(conn)
	if err != nil {
		log.Logf("[ss] %s - %s : %s", conn.RemoteAddr(), conn.LocalAddr(), err)
		return
	}
	// clear timer
	conn.SetReadDeadline(time.Time{})

	log.Logf("[ss] %s -> %s", conn.RemoteAddr(), addr)

	if !Can("tcp", addr, h.options.Whitelist, h.options.Blacklist) {
		log.Logf("[ss] Unauthorized to tcp connect to %s", addr)
		return
	}

	cc, err := h.options.Chain.Dial(addr)
	if err != nil {
		log.Logf("[ss] %s -> %s : %s", conn.RemoteAddr(), addr, err)
		return
	}
	defer cc.Close()

	log.Logf("[ss] %s <-> %s", conn.RemoteAddr(), addr)
	transport(conn, cc)
	log.Logf("[ss] %s >-< %s", conn.RemoteAddr(), addr)
}

const (
	idType  = 0 // address type index
	idIP0   = 1 // ip address start index
	idDmLen = 1 // domain address length index
	idDm0   = 2 // domain address start index

	typeIPv4 = 1 // type is ipv4 address
	typeDm   = 3 // type is domain address
	typeIPv6 = 4 // type is ipv6 address

	lenIPv4     = net.IPv4len + 2 // ipv4 + 2port
	lenIPv6     = net.IPv6len + 2 // ipv6 + 2port
	lenDmBase   = 2               // 1addrLen + 2port, plus addrLen
	lenHmacSha1 = 10
)

// This function is copied from shadowsocks library with some modification.
func (h *shadowHandler) getRequest(r io.Reader) (host string, err error) {
	// buf size should at least have the same size with the largest possible
	// request size (when addrType is 3, domain name has at most 256 bytes)
	// 1(addrType) + 1(lenByte) + 256(max length address) + 2(port)
	buf := make([]byte, smallBufferSize)

	// read till we get possible domain length field
	if _, err = io.ReadFull(r, buf[:idType+1]); err != nil {
		return
	}

	var reqStart, reqEnd int
	addrType := buf[idType]
	switch addrType & ss.AddrMask {
	case typeIPv4:
		reqStart, reqEnd = idIP0, idIP0+lenIPv4
	case typeIPv6:
		reqStart, reqEnd = idIP0, idIP0+lenIPv6
	case typeDm:
		if _, err = io.ReadFull(r, buf[idType+1:idDmLen+1]); err != nil {
			return
		}
		reqStart, reqEnd = idDm0, idDm0+int(buf[idDmLen])+lenDmBase
	default:
		err = fmt.Errorf("addr type %d not supported", addrType&ss.AddrMask)
		return
	}

	if _, err = io.ReadFull(r, buf[reqStart:reqEnd]); err != nil {
		return
	}

	// Return string for typeIP is not most efficient, but browsers (Chrome,
	// Safari, Firefox) all seems using typeDm exclusively. So this is not a
	// big problem.
	switch addrType & ss.AddrMask {
	case typeIPv4:
		host = net.IP(buf[idIP0 : idIP0+net.IPv4len]).String()
	case typeIPv6:
		host = net.IP(buf[idIP0 : idIP0+net.IPv6len]).String()
	case typeDm:
		host = string(buf[idDm0 : idDm0+int(buf[idDmLen])])
	}
	// parse port
	port := binary.BigEndian.Uint16(buf[reqEnd-2 : reqEnd])
	host = net.JoinHostPort(host, strconv.Itoa(int(port)))
	return
}

type shadowUDPListener struct {
	ln       net.PacketConn
	conns    map[string]*udpServerConn
	connChan chan net.Conn
	errChan  chan error
	ttl      time.Duration
}

// ShadowUDPListener creates a Listener for shadowsocks UDP relay server.
func ShadowUDPListener(addr string, cipher *url.Userinfo, ttl time.Duration) (Listener, error) {
	laddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	ln, err := net.ListenUDP("udp", laddr)
	if err != nil {
		return nil, err
	}

	var method, password string
	if cipher != nil {
		method = cipher.Username()
		password, _ = cipher.Password()
	}
	cp, err := ss.NewCipher(method, password)
	if err != nil {
		ln.Close()
		return nil, err
	}
	l := &shadowUDPListener{
		ln:       ss.NewSecurePacketConn(ln, cp, false),
		conns:    make(map[string]*udpServerConn),
		connChan: make(chan net.Conn, 1024),
		errChan:  make(chan error, 1),
		ttl:      ttl,
	}
	go l.listenLoop()
	return l, nil
}

func (l *shadowUDPListener) listenLoop() {
	for {
		b := make([]byte, mediumBufferSize)
		n, raddr, err := l.ln.ReadFrom(b)
		if err != nil {
			log.Logf("[ssu] peer -> %s : %s", l.Addr(), err)
			l.ln.Close()
			l.errChan <- err
			close(l.errChan)
			return
		}
		if Debug {
			log.Logf("[ssu] %s >>> %s : length %d", raddr, l.Addr(), n)
		}

		conn, ok := l.conns[raddr.String()]
		if !ok || conn.Closed() {
			conn = newUDPServerConn(l.ln, raddr, l.ttl)
			l.conns[raddr.String()] = conn

			select {
			case l.connChan <- conn:
			default:
				conn.Close()
				log.Logf("[ssu] %s - %s: connection queue is full", raddr, l.Addr())
			}
		}

		select {
		case conn.rChan <- b[:n]: // we keep the addr info so that the handler can identify the destination.
		default:
			log.Logf("[ssu] %s -> %s : read queue is full", raddr, l.Addr())
		}
	}
}

func (l *shadowUDPListener) Accept() (conn net.Conn, err error) {
	var ok bool
	select {
	case conn = <-l.connChan:
	case err, ok = <-l.errChan:
		if !ok {
			err = errors.New("accpet on closed listener")
		}
	}
	return
}

func (l *shadowUDPListener) Addr() net.Addr {
	return l.ln.LocalAddr()
}

func (l *shadowUDPListener) Close() error {
	return l.ln.Close()
}

type shadowUDPdHandler struct {
	ttl     time.Duration
	options *HandlerOptions
}

// ShadowUDPdHandler creates a server Handler for shadowsocks UDP relay server.
func ShadowUDPdHandler(opts ...HandlerOption) Handler {
	h := &shadowUDPdHandler{
		options: &HandlerOptions{},
	}
	for _, opt := range opts {
		opt(h.options)
	}
	return h
}

func (h *shadowUDPdHandler) Handle(conn net.Conn) {
	defer conn.Close()

	var err error
	var cc net.PacketConn
	if h.options.Chain.IsEmpty() {
		cc, err = net.ListenUDP("udp", nil)
		if err != nil {
			log.Logf("[ssu] %s - : %s", conn.LocalAddr(), err)
			return
		}
	} else {
		var c net.Conn
		c, err = getSOCKS5UDPTunnel(h.options.Chain, nil)
		if err != nil {
			log.Logf("[ssu] %s - : %s", conn.LocalAddr(), err)
			return
		}
		cc = &udpTunnelConn{Conn: c}
	}
	defer cc.Close()

	log.Logf("[ssu] %s <-> %s", conn.RemoteAddr(), conn.LocalAddr())
	transportUDP(conn, cc)
	log.Logf("[ssu] %s >-< %s", conn.RemoteAddr(), conn.LocalAddr())
}

func transportUDP(sc net.Conn, cc net.PacketConn) error {
	errc := make(chan error, 1)
	go func() {
		for {
			b := make([]byte, mediumBufferSize)
			n, err := sc.Read(b[3:]) // add rsv and frag fields to make it the standard SOCKS5 UDP datagram
			if err != nil {
				// log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), err)
				errc <- err
				return
			}
			dgram, err := gosocks5.ReadUDPDatagram(bytes.NewReader(b[:n+3]))
			if err != nil {
				log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), err)
				errc <- err
				return
			}
			//if Debug {
			//	log.Logf("[ssu] %s >>> %s length: %d", sc.RemoteAddr(), dgram.Header.Addr.String(), len(dgram.Data))
			//}
			addr, err := net.ResolveUDPAddr("udp", dgram.Header.Addr.String())
			if err != nil {
				errc <- err
				return
			}
			if _, err := cc.WriteTo(dgram.Data, addr); err != nil {
				errc <- err
				return
			}
		}
	}()

	go func() {
		for {
			b := make([]byte, mediumBufferSize)
			n, addr, err := cc.ReadFrom(b)
			if err != nil {
				errc <- err
				return
			}
			//if Debug {
			//	log.Logf("[ssu] %s <<< %s length: %d", sc.RemoteAddr(), addr, n)
			//}
			dgram := gosocks5.NewUDPDatagram(gosocks5.NewUDPHeader(0, 0, toSocksAddr(addr)), b[:n])
			buf := bytes.Buffer{}
			dgram.Write(&buf)
			if buf.Len() < 10 {
				log.Logf("[ssu] %s <- %s : invalid udp datagram", sc.RemoteAddr(), addr)
				continue
			}
			if _, err := sc.Write(buf.Bytes()[3:]); err != nil {
				errc <- err
				return
			}
		}
	}()

	err := <-errc
	if err != nil && err == io.EOF {
		err = nil
	}
	return err
}
<|endoftext|>


<commit_before>package redisence

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	gredis "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

type Status int

const (
	Offline Status = iota
	Online
	Closed
)

// Event is the data type for
// occuring events in the system
type Event struct {
	// Id is the given key by the application
	Id string

	// Status holds the changing type of event
	Status Status
}

// Prefix for redisence package
const RedisencePrefix = "redisence"

// Session holds the required connection data for redis
type Session struct {
	// main redis connection
	redis *redis.RedisSession

	// inactiveDuration specifies no-probe allowance time
	inactiveDuration time.Duration

	// receiving offline events pattern
	becameOfflinePattern string

	// receiving online events pattern
	becameOnlinePattern string
}

func New(server string, db int, inactiveDuration time.Duration) (*Session, error) {
	redis, err := redis.NewRedisSession(&redis.RedisConf{Server: server, DB: db})
	if err != nil {
		return nil, err
	}
	redis.SetPrefix(RedisencePrefix)

	return &Session{
		redis:                redis,
		becameOfflinePattern: fmt.Sprintf("__keyevent@%d__:expired", db),
		becameOnlinePattern:  fmt.Sprintf("__keyevent@%d__:set", db),
		inactiveDuration:     inactiveDuration,
	}, nil
}

// Ping resets the expiration time for any given key
// if key doesnt exists, it means user is now online
// Whenever application gets any prob from a client
// should call this function
func (s *Session) Ping(ids ...string) error {
	if len(ids) == 1 {
		// if member exits increase ttl
		if s.redis.Expire(ids[0], s.inactiveDuration) == nil {
			return nil
		}

		// if member doesnt exist set it
		return s.redis.Setex(ids[0], s.inactiveDuration, ids[0])
	}

	existance, err := s.sendMultiExpire(ids)
	if err != nil {
		return err
	}

	return s.sendMultiSetIfRequired(ids, existance)
}

func (s *Session) sendMultiSetIfRequired(ids []string, existance []int) error {
	if len(ids) != len(existance) {
		return fmt.Errorf("Length is not same Ids: %d Existance: %d", len(ids), len(existance))
	}

	// cache inactive duration as string
	seconds := strconv.Itoa(int(s.inactiveDuration.Seconds()))

	// get one connection from pool
	c := s.redis.Pool().Get()

	// item count for non-existent members
	notExistsCount := 0

	for i, exists := range existance {
		// `0` means, member doesnt exists in presence system
		if exists != 0 {
			continue
		}

		// init multi command lazily
		if notExistsCount == 0 {
			c.Send("MULTI")
		}

		notExistsCount++
		c.Send("SETEX", s.redis.AddPrefix(ids[i]), seconds, ids[i])
	}

	// execute multi command if only we flushed some to connection
	if notExistsCount != 0 {
		// ignore values
		if _, err := c.Do("EXEC"); err != nil {
			return err
		}
	}

	// do not forget to close the connection
	if err := c.Close(); err != nil {
		return err
	}

	return nil
}

func (s *Session) sendMultiExpire(ids []string) ([]int, error) {
	c := s.redis.Pool().Get()

	c.Send("MULTI")
	seconds := strconv.Itoa(int(s.inactiveDuration.Seconds()))
	for _, id := range ids {
		c.Send("EXPIRE", s.redis.AddPrefix(id), seconds)
	}
	r, err := c.Do("EXEC")
	if err != nil {
		return make([]int, 0), err
	}

	values, err := s.redis.Values(r)
	if err != nil {
		return make([]int, 0), err
	}

	err = c.Close()
	if err != nil {
		return make([]int, 0), err
	}

	res := make([]int, len(values))
	for i, value := range values {
		res[i], err = s.redis.Int(value)
		if err != nil {
			return make([]int, 0), err
		}

	}

	return res, nil
}

// Status returns the current status a key from system
// TODO use variadic function arguments
func (s *Session) Status(id string) Status {
	// to-do use MGET instead of exists
	if s.redis.Exists(id) {
		return Online
	}

	return Offline
}

// createEvent Creates the event with the required properties
func (s *Session) createEvent(n gredis.PMessage) Event {
	e := Event{}

	switch n.Pattern {
	case s.becameOfflinePattern:
		e.Id = string(n.Data[len(RedisencePrefix)+1:])
		e.Status = Offline
	case s.becameOnlinePattern:
		e.Id = string(n.Data[len(RedisencePrefix)+1:])
		e.Status = Online
	default:
		//ignore other events
	}

	return e
}

// ListenStatusChanges pubscribes to the redis and
// gets online and offline status changes from it
func (s *Session) ListenStatusChanges(events chan Event) {
	psc := s.redis.CreatePubSubConn()

	psc.PSubscribe(s.becameOnlinePattern, s.becameOfflinePattern)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		for {
			switch n := psc.Receive().(type) {
			case gredis.PMessage:
				events <- s.createEvent(n)
			case error:
				fmt.Printf("error: %v\n", n)
				return
			}
		}
	}()

	// avoid lock
	go func() {
		wg.Wait()
		psc.PUnsubscribe(s.becameOfflinePattern, s.becameOnlinePattern)
		psc.Close()
		events <- Event{Status: Closed}
	}()
}
<commit_msg>Redisence: close redis connection before processing the data<commit_after>package redisence

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	gredis "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

type Status int

const (
	Offline Status = iota
	Online
	Closed
)

// Event is the data type for
// occuring events in the system
type Event struct {
	// Id is the given key by the application
	Id string

	// Status holds the changing type of event
	Status Status
}

// Prefix for redisence package
const RedisencePrefix = "redisence"

// Session holds the required connection data for redis
type Session struct {
	// main redis connection
	redis *redis.RedisSession

	// inactiveDuration specifies no-probe allowance time
	inactiveDuration time.Duration

	// receiving offline events pattern
	becameOfflinePattern string

	// receiving online events pattern
	becameOnlinePattern string
}

func New(server string, db int, inactiveDuration time.Duration) (*Session, error) {
	redis, err := redis.NewRedisSession(&redis.RedisConf{Server: server, DB: db})
	if err != nil {
		return nil, err
	}
	redis.SetPrefix(RedisencePrefix)

	return &Session{
		redis:                redis,
		becameOfflinePattern: fmt.Sprintf("__keyevent@%d__:expired", db),
		becameOnlinePattern:  fmt.Sprintf("__keyevent@%d__:set", db),
		inactiveDuration:     inactiveDuration,
	}, nil
}

// Ping resets the expiration time for any given key
// if key doesnt exists, it means user is now online
// Whenever application gets any prob from a client
// should call this function
func (s *Session) Ping(ids ...string) error {
	if len(ids) == 1 {
		// if member exits increase ttl
		if s.redis.Expire(ids[0], s.inactiveDuration) == nil {
			return nil
		}

		// if member doesnt exist set it
		return s.redis.Setex(ids[0], s.inactiveDuration, ids[0])
	}

	existance, err := s.sendMultiExpire(ids)
	if err != nil {
		return err
	}

	return s.sendMultiSetIfRequired(ids, existance)
}

func (s *Session) sendMultiSetIfRequired(ids []string, existance []int) error {
	if len(ids) != len(existance) {
		return fmt.Errorf("Length is not same Ids: %d Existance: %d", len(ids), len(existance))
	}

	// cache inactive duration as string
	seconds := strconv.Itoa(int(s.inactiveDuration.Seconds()))

	// get one connection from pool
	c := s.redis.Pool().Get()

	// item count for non-existent members
	notExistsCount := 0

	for i, exists := range existance {
		// `0` means, member doesnt exists in presence system
		if exists != 0 {
			continue
		}

		// init multi command lazily
		if notExistsCount == 0 {
			c.Send("MULTI")
		}

		notExistsCount++
		c.Send("SETEX", s.redis.AddPrefix(ids[i]), seconds, ids[i])
	}

	// execute multi command if only we flushed some to connection
	if notExistsCount != 0 {
		// ignore values
		if _, err := c.Do("EXEC"); err != nil {
			return err
		}
	}

	// do not forget to close the connection
	if err := c.Close(); err != nil {
		return err
	}

	return nil
}

func (s *Session) sendMultiExpire(ids []string) ([]int, error) {
	// cache inactive duration as string
	seconds := strconv.Itoa(int(s.inactiveDuration.Seconds()))

	// get one connection from pool
	c := s.redis.Pool().Get()

	// init multi command
	c.Send("MULTI")

	// send expire command for all members
	for _, id := range ids {
		c.Send("EXPIRE", s.redis.AddPrefix(id), seconds)
	}

	// execute command
	r, err := c.Do("EXEC")
	if err != nil {
		return make([]int, 0), err
	}

	// close connection
	if err := c.Close(); err != nil {
		return make([]int, 0), err
	}

	values, err := s.redis.Values(r)
	if err != nil {
		return make([]int, 0), err
	}

	res := make([]int, len(values))
	for i, value := range values {
		res[i], err = s.redis.Int(value)
		if err != nil {
			return make([]int, 0), err
		}

	}

	return res, nil
}

// Status returns the current status a key from system
// TODO use variadic function arguments
func (s *Session) Status(id string) Status {
	// to-do use MGET instead of exists
	if s.redis.Exists(id) {
		return Online
	}

	return Offline
}

// createEvent Creates the event with the required properties
func (s *Session) createEvent(n gredis.PMessage) Event {
	e := Event{}

	switch n.Pattern {
	case s.becameOfflinePattern:
		e.Id = string(n.Data[len(RedisencePrefix)+1:])
		e.Status = Offline
	case s.becameOnlinePattern:
		e.Id = string(n.Data[len(RedisencePrefix)+1:])
		e.Status = Online
	default:
		//ignore other events
	}

	return e
}

// ListenStatusChanges pubscribes to the redis and
// gets online and offline status changes from it
func (s *Session) ListenStatusChanges(events chan Event) {
	psc := s.redis.CreatePubSubConn()

	psc.PSubscribe(s.becameOnlinePattern, s.becameOfflinePattern)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		for {
			switch n := psc.Receive().(type) {
			case gredis.PMessage:
				events <- s.createEvent(n)
			case error:
				fmt.Printf("error: %v\n", n)
				return
			}
		}
	}()

	// avoid lock
	go func() {
		wg.Wait()
		psc.PUnsubscribe(s.becameOfflinePattern, s.becameOnlinePattern)
		psc.Close()
		events <- Event{Status: Closed}
	}()
}
<|endoftext|>


<commit_before>package ds

import (
	"context"
	"reflect"

	"cloud.google.com/go/datastore"
)

// Tx is the datastore transaction wrapper
type Tx struct {
	*datastore.Transaction
	invalidateKeys []*datastore.Key
}

// RunInTx is the RunInTransaction wrapper
func (client *Client) RunInTx(ctx context.Context, f func(tx *Tx) error, opts ...datastore.TransactionOption) (*datastore.Commit, error) {
	var tx *Tx
	commit, err := client.RunInTransaction(ctx, func(t *datastore.Transaction) error {
		tx = &Tx{t, nil}
		return f(tx)
	})
	if err == nil && client.Cache != nil && len(tx.invalidateKeys) > 0 {
		// find unique keys
		mapKey := map[string]*datastore.Key{}
		for _, key := range tx.invalidateKeys {
			mapKey[key.String()] = key
		}
		uKeys := []*datastore.Key{}
		for _, k := range mapKey {
			uKeys = append(uKeys, k)
		}
		client.Cache.DelMulti(uKeys)
	}

	return commit, err
}

// GetByKey retrieves model from datastore by key
func (tx *Tx) GetByKey(key *datastore.Key, dst interface{}) error {
	err := tx.Get(key, dst)
	SetKey(key, dst)
	if err != nil {
		return err
	}
	return nil
}

// GetByKeys retrieves models from datastore by keys
func (tx *Tx) GetByKeys(keys []*datastore.Key, dst interface{}) error {
	// prepare slice if dst is pointer to 0 len slice
	if rf := reflect.ValueOf(dst); rf.Kind() == reflect.Ptr {
		rs := rf.Elem()
		if rs.Kind() == reflect.Slice && rs.Len() == 0 {
			l := len(keys)
			rs.Set(reflect.MakeSlice(rs.Type(), l, l))
		}
		dst = rs.Interface()
	}

	err := tx.GetMulti(keys, dst)
	SetKeys(keys, dst)
	if err != nil {
		return err
	}
	return nil
}

// GetByModel retrieves model from datastore by key from model
func (tx *Tx) GetByModel(dst interface{}) error {
	key := ExtractKey(dst)
	return tx.GetByKey(key, dst)
}

// GetByModels retrieves models from datastore by keys from models
func (tx *Tx) GetByModels(dst interface{}) error {
	keys := ExtractKeys(dst)
	return tx.GetByKeys(keys, dst)
}

// GetByID retrieves model from datastore by id
func (tx *Tx) GetByID(kind string, id int64, dst interface{}) error {
	return tx.GetByKey(datastore.IDKey(kind, id, nil), dst)
}

// GetByIDs retrieves models from datastore by ids
func (tx *Tx) GetByIDs(kind string, ids []int64, dst interface{}) error {
	keys := BuildIDKeys(kind, ids)
	return tx.GetByKeys(keys, dst)
}

// GetByStringID retrieves model from datastore by string id
func (tx *Tx) GetByStringID(kind string, id string, dst interface{}) error {
	tid := parseID(id)
	if tid == 0 {
		return datastore.ErrInvalidKey
	}
	return tx.GetByKey(datastore.IDKey(kind, tid, nil), dst)
}

// GetByStringIDs retrieves models from datastore by string ids
func (tx *Tx) GetByStringIDs(kind string, ids []string, dst interface{}) error {
	keys := BuildStringIDKeys(kind, ids)
	return tx.GetByKeys(keys, dst)
}

// GetByName retrieves model from datastore by name
func (tx *Tx) GetByName(kind string, name string, dst interface{}) error {
	return tx.GetByKey(datastore.NameKey(kind, name, nil), dst)
}

// GetByNames retrieves models from datastore by names
func (tx *Tx) GetByNames(kind string, names []string, dst interface{}) error {
	keys := BuildNameKeys(kind, names)
	return tx.GetByKeys(keys, dst)
}

// PutModel puts a model to datastore
func (tx *Tx) PutModel(src interface{}) (*datastore.PendingKey, error) {
	x := src.(KeyGetSetter)
	key := x.GetKey()
	pKey, err := tx.Put(key, x)
	if key != nil && !key.Incomplete() {
		tx.invalidateKeys = append(tx.invalidateKeys, key)
	}
	return pKey, err
}

// PutModels puts models to datastore
func (tx *Tx) PutModels(src interface{}) ([]*datastore.PendingKey, error) {
	xs := valueOf(src)
	keys := make([]*datastore.Key, xs.Len())
	for i := range keys {
		x := xs.Index(i).Interface()
		keys[i] = x.(KeyGetter).GetKey()
	}
	// TODO: store pending key inside model ?
	ks, err := tx.PutMulti(keys, src)
	for _, key := range keys {
		if key != nil && !key.Incomplete() {
			tx.invalidateKeys = append(tx.invalidateKeys, key)
		}
	}
	return ks, err
}

// SaveModel saves model to datastore
func (tx *Tx) SaveModel(src interface{}) (*datastore.PendingKey, error) {
	beforeSave(src)
	return tx.PutModel(src)
}

// SaveModels saves models to datastore
func (tx *Tx) SaveModels(src interface{}) ([]*datastore.PendingKey, error) {
	xs := valueOf(src)
	for i := 0; i < xs.Len(); i++ {
		x := xs.Index(i).Interface()
		beforeSave(x)
	}
	return tx.PutModels(src)
}
<commit_msg>add delete by key for tx<commit_after>package ds

import (
	"context"
	"reflect"

	"cloud.google.com/go/datastore"
)

// Tx is the datastore transaction wrapper
type Tx struct {
	*datastore.Transaction
	invalidateKeys []*datastore.Key
}

// RunInTx is the RunInTransaction wrapper
func (client *Client) RunInTx(ctx context.Context, f func(tx *Tx) error, opts ...datastore.TransactionOption) (*datastore.Commit, error) {
	var tx *Tx
	commit, err := client.RunInTransaction(ctx, func(t *datastore.Transaction) error {
		tx = &Tx{t, nil}
		return f(tx)
	})
	if err == nil && client.Cache != nil && len(tx.invalidateKeys) > 0 {
		// find unique keys
		mapKey := map[string]*datastore.Key{}
		for _, key := range tx.invalidateKeys {
			mapKey[key.String()] = key
		}
		uKeys := []*datastore.Key{}
		for _, k := range mapKey {
			uKeys = append(uKeys, k)
		}
		client.Cache.DelMulti(uKeys)
	}

	return commit, err
}

// GetByKey retrieves model from datastore by key
func (tx *Tx) GetByKey(key *datastore.Key, dst interface{}) error {
	err := tx.Get(key, dst)
	SetKey(key, dst)
	if err != nil {
		return err
	}
	return nil
}

// GetByKeys retrieves models from datastore by keys
func (tx *Tx) GetByKeys(keys []*datastore.Key, dst interface{}) error {
	// prepare slice if dst is pointer to 0 len slice
	if rf := reflect.ValueOf(dst); rf.Kind() == reflect.Ptr {
		rs := rf.Elem()
		if rs.Kind() == reflect.Slice && rs.Len() == 0 {
			l := len(keys)
			rs.Set(reflect.MakeSlice(rs.Type(), l, l))
		}
		dst = rs.Interface()
	}

	err := tx.GetMulti(keys, dst)
	SetKeys(keys, dst)
	if err != nil {
		return err
	}
	return nil
}

// GetByModel retrieves model from datastore by key from model
func (tx *Tx) GetByModel(dst interface{}) error {
	key := ExtractKey(dst)
	return tx.GetByKey(key, dst)
}

// GetByModels retrieves models from datastore by keys from models
func (tx *Tx) GetByModels(dst interface{}) error {
	keys := ExtractKeys(dst)
	return tx.GetByKeys(keys, dst)
}

// GetByID retrieves model from datastore by id
func (tx *Tx) GetByID(kind string, id int64, dst interface{}) error {
	return tx.GetByKey(datastore.IDKey(kind, id, nil), dst)
}

// GetByIDs retrieves models from datastore by ids
func (tx *Tx) GetByIDs(kind string, ids []int64, dst interface{}) error {
	keys := BuildIDKeys(kind, ids)
	return tx.GetByKeys(keys, dst)
}

// GetByStringID retrieves model from datastore by string id
func (tx *Tx) GetByStringID(kind string, id string, dst interface{}) error {
	tid := parseID(id)
	if tid == 0 {
		return datastore.ErrInvalidKey
	}
	return tx.GetByKey(datastore.IDKey(kind, tid, nil), dst)
}

// GetByStringIDs retrieves models from datastore by string ids
func (tx *Tx) GetByStringIDs(kind string, ids []string, dst interface{}) error {
	keys := BuildStringIDKeys(kind, ids)
	return tx.GetByKeys(keys, dst)
}

// GetByName retrieves model from datastore by name
func (tx *Tx) GetByName(kind string, name string, dst interface{}) error {
	return tx.GetByKey(datastore.NameKey(kind, name, nil), dst)
}

// GetByNames retrieves models from datastore by names
func (tx *Tx) GetByNames(kind string, names []string, dst interface{}) error {
	keys := BuildNameKeys(kind, names)
	return tx.GetByKeys(keys, dst)
}

// PutModel puts a model to datastore
func (tx *Tx) PutModel(src interface{}) (*datastore.PendingKey, error) {
	x := src.(KeyGetSetter)
	key := x.GetKey()
	pKey, err := tx.Put(key, x)
	if key != nil && !key.Incomplete() {
		tx.invalidateKeys = append(tx.invalidateKeys, key)
	}
	return pKey, err
}

// PutModels puts models to datastore
func (tx *Tx) PutModels(src interface{}) ([]*datastore.PendingKey, error) {
	xs := valueOf(src)
	keys := make([]*datastore.Key, xs.Len())
	for i := range keys {
		x := xs.Index(i).Interface()
		keys[i] = x.(KeyGetter).GetKey()
	}
	// TODO: store pending key inside model ?
	ks, err := tx.PutMulti(keys, src)
	for _, key := range keys {
		if key != nil && !key.Incomplete() {
			tx.invalidateKeys = append(tx.invalidateKeys, key)
		}
	}
	return ks, err
}

// SaveModel saves model to datastore
func (tx *Tx) SaveModel(src interface{}) (*datastore.PendingKey, error) {
	beforeSave(src)
	return tx.PutModel(src)
}

// SaveModels saves models to datastore
func (tx *Tx) SaveModels(src interface{}) ([]*datastore.PendingKey, error) {
	xs := valueOf(src)
	for i := 0; i < xs.Len(); i++ {
		x := xs.Index(i).Interface()
		beforeSave(x)
	}
	return tx.PutModels(src)
}

// DeleteByKey deletes a model by key
func (tx *Tx) DeleteByKey(key *datastore.Key) error {
	err := tx.Delete(key)
	if key != nil && !key.Incomplete() {
		tx.invalidateKeys = append(tx.invalidateKeys, key)
	}
	return err
}

// DeleteByKeys deletes models by keys
func (tx *Tx) DeleteByKeys(keys []*datastore.Key) error {
	err := tx.DeleteMulti(keys)
	for _, key := range keys {
		if key != nil && !key.Incomplete() {
			tx.invalidateKeys = append(tx.invalidateKeys, key)
		}
	}
	return err
}
<|endoftext|>


<commit_before>package docker

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"syscall"
	"time"

	"google.golang.org/cloud/datastore"
)

// Config holds the configuration needed for a docker run
type Config struct {
	Image,
	Volume,
	Code,
	User string
	Challenge *datastore.Key
}

// Rusage is is a copy of syscall.Rusage for Linux.
// It is needed as syscall.Rusage built on Windows cannot be saved in Datastore
// because it holds unsigned integers. Runs are executed inside Docker and
// therefore will always generate this version of syscall.Rusage.
// See https://godoc.org/google.golang.org/cloud/datastore#Property
// See https://golang.org/src/syscall/syscall_windows.go
// See https://golang.org/src/syscall/ztypes_linux_amd64.go
type Rusage struct {
	Utime,
	Stime syscall.Timeval
	Maxrss,
	Ixrss,
	Idrss,
	Isrss,
	Minflt,
	Majflt,
	Nswap,
	Inblock,
	Oublock,
	Msgsnd,
	Msgrcv,
	Nsignals,
	Nvcsw,
	Nivcsw int64
}

const volumePattern string = "coduno-volume"
const imagePattern string = "coduno/fingerprint-"

// NewConfig constructs a Config to run the specfied image
// using the given volume.
// If volume is left blank, NewConfig will take care of
// creating a temporary directory.
func NewConfig(image, volume, code string) (c *Config, err error) {
	if volume == "" {
		volume, err = volumeDir()
	}
	c = &Config{
		Image:  image,
		Volume: volume,
		Code:   code,
	}
	return
}

// Result holds the results (standard output and standard error)
// of a Config that was run.
type Result struct {
	Config
	Stdout     string
	Stderr     string
	Rusage     Rusage
	Prepare    string
	Exit       error
	Start, End time.Time
}

// Run executes a Config and returns associated results.
func (c *Config) Run() (r *Result, err error) {
	r = new(Result)
	dockerized, err := dockerize(c.Volume)

	if err != nil {
		return
	}

	cmd := exec.Command(
		"docker",
		"run",
		"-v",
		dockerized+":/run",
		c.Image,
	)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return
	}
	r.Start = time.Now()
	if err = cmd.Start(); err != nil {
		return
	}

	bufout, buferr := new(bytes.Buffer), new(bytes.Buffer)

	go io.Copy(io.MultiWriter(os.Stdout, bufout), stdout)
	go io.Copy(io.MultiWriter(os.Stdout, buferr), stderr)

	r.Exit = cmd.Wait()
	r.End = time.Now()
	r.Stdout = bufout.String()
	r.Stderr = buferr.String()
	b, err := ioutil.ReadFile(path.Join(c.Volume, "prepare.log"))
	r.Prepare = string(b)

	stats, err := os.Open(path.Join(c.Volume, "stats.log"))
	json.NewDecoder(stats).Decode(&r.Rusage)
	return
}

// NewImage returns the correct docker image name for a
// specific language.
func NewImage(language string) string {
	return imagePattern + language
}
<commit_msg>docker: Make results contain exit error as string<commit_after>package docker

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"syscall"
	"time"

	"google.golang.org/cloud/datastore"
)

// Config holds the configuration needed for a docker run
type Config struct {
	Image,
	Volume,
	Code,
	User string
	Challenge *datastore.Key
}

// Rusage is is a copy of syscall.Rusage for Linux.
// It is needed as syscall.Rusage built on Windows cannot be saved in Datastore
// because it holds unsigned integers. Runs are executed inside Docker and
// therefore will always generate this version of syscall.Rusage.
// See https://godoc.org/google.golang.org/cloud/datastore#Property
// See https://golang.org/src/syscall/syscall_windows.go
// See https://golang.org/src/syscall/ztypes_linux_amd64.go
type Rusage struct {
	Utime,
	Stime syscall.Timeval
	Maxrss,
	Ixrss,
	Idrss,
	Isrss,
	Minflt,
	Majflt,
	Nswap,
	Inblock,
	Oublock,
	Msgsnd,
	Msgrcv,
	Nsignals,
	Nvcsw,
	Nivcsw int64
}

const volumePattern string = "coduno-volume"
const imagePattern string = "coduno/fingerprint-"

// NewConfig constructs a Config to run the specfied image
// using the given volume.
// If volume is left blank, NewConfig will take care of
// creating a temporary directory.
func NewConfig(image, volume, code string) (c *Config, err error) {
	if volume == "" {
		volume, err = volumeDir()
	}
	c = &Config{
		Image:  image,
		Volume: volume,
		Code:   code,
	}
	return
}

// Result holds the results (standard output and standard error)
// of a Config that was run.
type Result struct {
	Config
	Stdout     string
	Stderr     string
	Rusage     Rusage
	Prepare    string
	Exit       string
	Start, End time.Time
}

// Run executes a Config and returns associated results.
func (c *Config) Run() (r *Result, err error) {
	r = new(Result)
	dockerized, err := dockerize(c.Volume)

	if err != nil {
		return
	}

	cmd := exec.Command(
		"docker",
		"run",
		"-v",
		dockerized+":/run",
		c.Image,
	)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return
	}
	r.Start = time.Now()
	if err = cmd.Start(); err != nil {
		return
	}

	bufout, buferr := new(bytes.Buffer), new(bytes.Buffer)

	go io.Copy(io.MultiWriter(os.Stdout, bufout), stdout)
	go io.Copy(io.MultiWriter(os.Stdout, buferr), stderr)

	err = cmd.Wait()
	if err != nil {
		r.Exit = err.Error()
	}
	r.End = time.Now()
	r.Stdout = bufout.String()
	r.Stderr = buferr.String()
	b, err := ioutil.ReadFile(path.Join(c.Volume, "prepare.log"))
	r.Prepare = string(b)

	stats, err := os.Open(path.Join(c.Volume, "stats.log"))
	json.NewDecoder(stats).Decode(&r.Rusage)
	return
}

// NewImage returns the correct docker image name for a
// specific language.
func NewImage(language string) string {
	return imagePattern + language
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"io"
	"log"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"time"
)

// Use direct connection after blocked for chouTimeout
const chouTimeout = 2 * time.Minute

type domainSet map[string]bool

// Basically a concurrent map. I don't want to use channels to implement
// concurrent access to this as I'm comfortable to use locks for simple tasks
// like this
type paraDomainSet struct {
	sync.RWMutex
	domainSet
}

func newDomainSet() domainSet {
	return make(map[string]bool)
}

func (ds domainSet) loadDomainList(fpath string) (lst []string, err error) {
	lst, err = loadDomainList(fpath)
	if err != nil {
		return
	}
	// This executes in single goroutine, so no need to use lock
	for _, v := range lst {
		// debug.Println("loaded domain:", v)
		ds[v] = true
	}
	return
}

func (ds domainSet) toArray() []string {
	l := len(ds)
	lst := make([]string, l, l)

	i := 0
	for k, _ := range ds {
		lst[i] = k
		i++
	}
	return lst
}

func newParaDomainSet() *paraDomainSet {
	return &paraDomainSet{domainSet: newDomainSet()}
}

func (ds *paraDomainSet) add(dm string) {
	ds.Lock()
	ds.domainSet[dm] = true
	ds.Unlock()
}

func (ds *paraDomainSet) has(dm string) bool {
	ds.RLock()
	_, ok := ds.domainSet[dm]
	ds.RUnlock()
	return ok
}

func (ds *paraDomainSet) del(dm string) {
	ds.Lock()
	delete(ds.domainSet, dm)
	ds.Unlock()
}

var blockedDs = newParaDomainSet()
var directDs = newParaDomainSet()

var blockedDomainChanged = false
var directDomainChanged = false

var alwaysBlockedDs = newDomainSet()
var alwaysDirectDs = newDomainSet()
var chouDs = newDomainSet()

// Record when is the domain added to chou domain set
type chouBlockTime struct {
	sync.Mutex
	time map[string]time.Time
}

var chou = chouBlockTime{time: map[string]time.Time{}}

func isHostInAlwaysDs(host string) bool {
	dm := host2Domain(host)
	return alwaysBlockedDs[dm] || alwaysDirectDs[dm]
}

func isHostAlwaysDirect(host string) bool {
	return alwaysDirectDs[host2Domain(host)]
}

func isHostAlwaysBlocked(host string) bool {
	return alwaysBlockedDs[host2Domain(host)]
}

func isHostBlocked(host string) bool {
	dm := host2Domain(host)
	if alwaysDirectDs[dm] {
		return false
	}
	if alwaysBlockedDs[dm] {
		return true
	}
	if chouDs[dm] {
		chou.Lock()
		t, ok := chou.time[dm]
		chou.Unlock()
		if !ok {
			return false
		}
		if time.Now().Sub(t) < chouTimeout {
			return true
		}
		chou.Lock()
		delete(chou.time, dm)
		chou.Unlock()
		debug.Printf("chou domain %s block time unset\n", dm)
		return false
	}
	return blockedDs.has(dm)
}

func isHostDirect(host string) bool {
	dm := host2Domain(host)
	if alwaysDirectDs[dm] {
		return true
	}
	if alwaysBlockedDs[dm] {
		return false
	}
	return directDs.has(dm)
}

func isHostChouFeng(host string) bool {
	return chouDs[host2Domain(host)]
}

// Return true if the host is taken as blocked later
func addBlockedHost(host string) bool {
	dm := host2Domain(host)
	if isHostAlwaysDirect(host) || hostIsIP(host) || dm == "localhost" {
		return false
	}
	if chouDs[dm] {
		// Record blocked time for chou domain, this marks a chou domain as
		// temporarily blocked
		now := time.Now()
		chou.Lock()
		chou.time[dm] = now
		chou.Unlock()
		debug.Printf("chou domain %s blocked at %v\n", dm, now)
	} else if !blockedDs.has(dm) {
		blockedDs.add(dm)
		blockedDomainChanged = true
		debug.Printf("%s added to blocked list\n", dm)
		// Delete this domain from direct domain set
		delDirectDomain(dm)
	}
	return true
}

func delBlockedDomain(dm string) {
	if blockedDs.has(dm) {
		blockedDs.del(dm)
		blockedDomainChanged = true
		debug.Printf("%s deleted from blocked list\n", dm)
	}
}

func addDirectDomain(dm string) {
	if !config.updateDirect {
		return
	}
}

func addDirectHost(host string) (added bool) {
	dm := host2Domain(host)
	if !config.updateDirect || isHostInAlwaysDs(host) || chouDs[dm] ||
		dm == "localhost" || hostIsIP(host) {
		return
	}
	if !directDs.has(dm) {
		directDs.add(dm)
		directDomainChanged = true
		debug.Printf("%s added to direct list\n", dm)
	}
	// Delete this domain from blocked domain set
	delBlockedDomain(dm)
	return true
}

func delDirectDomain(dm string) {
	if directDs.has(dm) {
		directDs.del(dm)
		directDomainChanged = true
	}
}

func writeBlockedDs() {
	if !config.updateBlocked || !blockedDomainChanged {
		return
	}
	writeDomainList(config.blockedFile, blockedDs.toArray())
}

func writeDirectDs() {
	if !config.updateDirect || !directDomainChanged {
		return
	}
	writeDomainList(config.directFile, directDs.toArray())
}

// filter out domain in blocked and direct domain set.
func filterOutDs(ds domainSet) {
	for k, _ := range ds {
		if blockedDs.domainSet[k] {
			delete(blockedDs.domainSet, k)
			blockedDomainChanged = true
		}
		if directDs.domainSet[k] {
			delete(directDs.domainSet, k)
			directDomainChanged = true
		}
	}
}

// If a domain name appears in both blocked and direct domain set, only keep
// it in the blocked set.
func filterOutBlockedDsInDirectDs() {
	for k, _ := range blockedDs.domainSet {
		if directDs.domainSet[k] {
			delete(directDs.domainSet, k)
			directDomainChanged = true
		}
	}
	for k, _ := range alwaysBlockedDs {
		if alwaysDirectDs[k] {
			errl.Printf("%s in both always blocked and direct domain lists, taken as blocked.\n", k)
			delete(alwaysDirectDs, k)
		}
	}
}

func writeDomainSet() {
	// chou domain set maybe added to blocked site during execution,
	// filter them out before writing back to disk.
	filterOutDs(chouDs)

	writeBlockedDs()
	writeDirectDs()
}

func loadDomainList(fpath string) (lst []string, err error) {
	f, err := openFile(fpath)
	if f == nil || err != nil {
		return
	}
	defer f.Close()

	fr := bufio.NewReader(f)
	lst = make([]string, 0)
	var domain string
	for {
		domain, err = ReadLine(fr)
		if err == io.EOF {
			return lst, nil
		} else if err != nil {
			errl.Println("Error reading domain list from:", fpath, err)
			return
		}
		if domain == "" {
			continue
		}
		lst = append(lst, strings.TrimSpace(domain))
	}
	return
}

func mkConfigDir() (err error) {
	stat, err := os.Stat(config.dir)
	if err == nil {
		if stat.IsDir() {
			return
		}
		log.Printf("%s is not directory, can't write domain list\n", config.dir)
		return errors.New("config.dir is not directory")
	}
	if os.IsNotExist(err) {
		err = os.Mkdir(config.dir, 0755)
	}
	if err != nil {
		log.Printf("Config directory %s: %v\n", config.dir, err)
	}
	return
}

func writeDomainList(fpath string, lst []string) (err error) {
	if err = mkConfigDir(); err != nil {
		return
	}
	tmpPath := path.Join(config.dir, "tmp-domain")
	f, err := os.Create(tmpPath)
	if err != nil {
		errl.Println("Error creating tmp domain list file:", err)
		return
	}

	sort.Sort(sort.StringSlice(lst))

	all := strings.Join(lst, "\n")
	f.WriteString(all)
	f.Close()

	// On windows, can't rename to a file which already exists.
	if isWindows() {
		if err = os.Remove(fpath); err != nil {
			errl.Println("Can't remove domain list", fpath, "for update")
		}
	}
	if err = os.Rename(tmpPath, fpath); err != nil {
		errl.Printf("Error moving tmp domain list file to %s: %v\n", fpath, err)
	}
	return
}

var topLevelDomain = map[string]bool{
	"ac":  true,
	"co":  true,
	"org": true,
	"com": true,
	"net": true,
	"edu": true,
}

func host2Domain(host string) (domain string) {
	host, _ = splitHostPort(host)
	host = trimLastDot(host)
	lastDot := strings.LastIndex(host, ".")
	if lastDot == -1 {
		return host // simple host name, we should not hanlde this
	}
	// Find the 2nd last dot
	dot2ndLast := strings.LastIndex(host[:lastDot], ".")
	if dot2ndLast == -1 {
		return host
	}

	part := host[dot2ndLast+1 : lastDot]
	// If the 2nd last part of a domain name equals to a top level
	// domain, search for the 3rd part in the host name.
	// So domains like bbc.co.uk will not be recorded as co.uk
	if topLevelDomain[part] {
		dot3rdLast := strings.LastIndex(host[:dot2ndLast], ".")
		if dot3rdLast == -1 {
			return host
		}
		return host[dot3rdLast+1:]
	}
	return host[dot2ndLast+1:]
}

// TODO reload domain list when receives SIGUSR1
// one difficult here is that we may concurrently access maps which is not
// safe.
// Can we create a new domain set first, then change the reference of the original one?
// Domain set reference changing should be atomic.

func loadDomainSet() {
	blockedDs.loadDomainList(config.blockedFile)
	directDs.loadDomainList(config.directFile)
	alwaysBlockedDs.loadDomainList(config.alwaysBlockedFile)
	alwaysDirectDs.loadDomainList(config.alwaysDirectFile)
	chouDs.loadDomainList(config.chouFile)

	filterOutDs(chouDs)
	filterOutDs(alwaysDirectDs)
	filterOutDs(alwaysBlockedDs)
	filterOutBlockedDsInDirectDs()
}
<commit_msg>Remove unused function addDirectDomain.<commit_after>package main

import (
	"bufio"
	"errors"
	"io"
	"log"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"time"
)

// Use direct connection after blocked for chouTimeout
const chouTimeout = 2 * time.Minute

type domainSet map[string]bool

// Basically a concurrent map. I don't want to use channels to implement
// concurrent access to this as I'm comfortable to use locks for simple tasks
// like this
type paraDomainSet struct {
	sync.RWMutex
	domainSet
}

func newDomainSet() domainSet {
	return make(map[string]bool)
}

func (ds domainSet) loadDomainList(fpath string) (lst []string, err error) {
	lst, err = loadDomainList(fpath)
	if err != nil {
		return
	}
	// This executes in single goroutine, so no need to use lock
	for _, v := range lst {
		// debug.Println("loaded domain:", v)
		ds[v] = true
	}
	return
}

func (ds domainSet) toArray() []string {
	l := len(ds)
	lst := make([]string, l, l)

	i := 0
	for k, _ := range ds {
		lst[i] = k
		i++
	}
	return lst
}

func newParaDomainSet() *paraDomainSet {
	return &paraDomainSet{domainSet: newDomainSet()}
}

func (ds *paraDomainSet) add(dm string) {
	ds.Lock()
	ds.domainSet[dm] = true
	ds.Unlock()
}

func (ds *paraDomainSet) has(dm string) bool {
	ds.RLock()
	_, ok := ds.domainSet[dm]
	ds.RUnlock()
	return ok
}

func (ds *paraDomainSet) del(dm string) {
	ds.Lock()
	delete(ds.domainSet, dm)
	ds.Unlock()
}

var blockedDs = newParaDomainSet()
var directDs = newParaDomainSet()

var blockedDomainChanged = false
var directDomainChanged = false

var alwaysBlockedDs = newDomainSet()
var alwaysDirectDs = newDomainSet()
var chouDs = newDomainSet()

// Record when is the domain added to chou domain set
type chouBlockTime struct {
	sync.Mutex
	time map[string]time.Time
}

var chou = chouBlockTime{time: map[string]time.Time{}}

func isHostInAlwaysDs(host string) bool {
	dm := host2Domain(host)
	return alwaysBlockedDs[dm] || alwaysDirectDs[dm]
}

func isHostAlwaysDirect(host string) bool {
	return alwaysDirectDs[host2Domain(host)]
}

func isHostAlwaysBlocked(host string) bool {
	return alwaysBlockedDs[host2Domain(host)]
}

func isHostBlocked(host string) bool {
	dm := host2Domain(host)
	if alwaysDirectDs[dm] {
		return false
	}
	if alwaysBlockedDs[dm] {
		return true
	}
	if chouDs[dm] {
		chou.Lock()
		t, ok := chou.time[dm]
		chou.Unlock()
		if !ok {
			return false
		}
		if time.Now().Sub(t) < chouTimeout {
			return true
		}
		chou.Lock()
		delete(chou.time, dm)
		chou.Unlock()
		debug.Printf("chou domain %s block time unset\n", dm)
		return false
	}
	return blockedDs.has(dm)
}

func isHostDirect(host string) bool {
	dm := host2Domain(host)
	if alwaysDirectDs[dm] {
		return true
	}
	if alwaysBlockedDs[dm] {
		return false
	}
	return directDs.has(dm)
}

func isHostChouFeng(host string) bool {
	return chouDs[host2Domain(host)]
}

// Return true if the host is taken as blocked later
func addBlockedHost(host string) bool {
	dm := host2Domain(host)
	if isHostAlwaysDirect(host) || hostIsIP(host) || dm == "localhost" {
		return false
	}
	if chouDs[dm] {
		// Record blocked time for chou domain, this marks a chou domain as
		// temporarily blocked
		now := time.Now()
		chou.Lock()
		chou.time[dm] = now
		chou.Unlock()
		debug.Printf("chou domain %s blocked at %v\n", dm, now)
	} else if !blockedDs.has(dm) {
		blockedDs.add(dm)
		blockedDomainChanged = true
		debug.Printf("%s added to blocked list\n", dm)
		// Delete this domain from direct domain set
		delDirectDomain(dm)
	}
	return true
}

func delBlockedDomain(dm string) {
	if blockedDs.has(dm) {
		blockedDs.del(dm)
		blockedDomainChanged = true
		debug.Printf("%s deleted from blocked list\n", dm)
	}
}

func addDirectHost(host string) (added bool) {
	dm := host2Domain(host)
	if !config.updateDirect || isHostInAlwaysDs(host) || chouDs[dm] ||
		dm == "localhost" || hostIsIP(host) {
		return
	}
	if !directDs.has(dm) {
		directDs.add(dm)
		directDomainChanged = true
		debug.Printf("%s added to direct list\n", dm)
	}
	// Delete this domain from blocked domain set
	delBlockedDomain(dm)
	return true
}

func delDirectDomain(dm string) {
	if directDs.has(dm) {
		directDs.del(dm)
		directDomainChanged = true
	}
}

func writeBlockedDs() {
	if !config.updateBlocked || !blockedDomainChanged {
		return
	}
	writeDomainList(config.blockedFile, blockedDs.toArray())
}

func writeDirectDs() {
	if !config.updateDirect || !directDomainChanged {
		return
	}
	writeDomainList(config.directFile, directDs.toArray())
}

// filter out domain in blocked and direct domain set.
func filterOutDs(ds domainSet) {
	for k, _ := range ds {
		if blockedDs.domainSet[k] {
			delete(blockedDs.domainSet, k)
			blockedDomainChanged = true
		}
		if directDs.domainSet[k] {
			delete(directDs.domainSet, k)
			directDomainChanged = true
		}
	}
}

// If a domain name appears in both blocked and direct domain set, only keep
// it in the blocked set.
func filterOutBlockedDsInDirectDs() {
	for k, _ := range blockedDs.domainSet {
		if directDs.domainSet[k] {
			delete(directDs.domainSet, k)
			directDomainChanged = true
		}
	}
	for k, _ := range alwaysBlockedDs {
		if alwaysDirectDs[k] {
			errl.Printf("%s in both always blocked and direct domain lists, taken as blocked.\n", k)
			delete(alwaysDirectDs, k)
		}
	}
}

func writeDomainSet() {
	// chou domain set maybe added to blocked site during execution,
	// filter them out before writing back to disk.
	filterOutDs(chouDs)

	writeBlockedDs()
	writeDirectDs()
}

func loadDomainList(fpath string) (lst []string, err error) {
	f, err := openFile(fpath)
	if f == nil || err != nil {
		return
	}
	defer f.Close()

	fr := bufio.NewReader(f)
	lst = make([]string, 0)
	var domain string
	for {
		domain, err = ReadLine(fr)
		if err == io.EOF {
			return lst, nil
		} else if err != nil {
			errl.Println("Error reading domain list from:", fpath, err)
			return
		}
		if domain == "" {
			continue
		}
		lst = append(lst, strings.TrimSpace(domain))
	}
	return
}

func mkConfigDir() (err error) {
	stat, err := os.Stat(config.dir)
	if err == nil {
		if stat.IsDir() {
			return
		}
		log.Printf("%s is not directory, can't write domain list\n", config.dir)
		return errors.New("config.dir is not directory")
	}
	if os.IsNotExist(err) {
		err = os.Mkdir(config.dir, 0755)
	}
	if err != nil {
		log.Printf("Config directory %s: %v\n", config.dir, err)
	}
	return
}

func writeDomainList(fpath string, lst []string) (err error) {
	if err = mkConfigDir(); err != nil {
		return
	}
	tmpPath := path.Join(config.dir, "tmp-domain")
	f, err := os.Create(tmpPath)
	if err != nil {
		errl.Println("Error creating tmp domain list file:", err)
		return
	}

	sort.Sort(sort.StringSlice(lst))

	all := strings.Join(lst, "\n")
	f.WriteString(all)
	f.Close()

	// On windows, can't rename to a file which already exists.
	if isWindows() {
		if err = os.Remove(fpath); err != nil {
			errl.Println("Can't remove domain list", fpath, "for update")
		}
	}
	if err = os.Rename(tmpPath, fpath); err != nil {
		errl.Printf("Error moving tmp domain list file to %s: %v\n", fpath, err)
	}
	return
}

var topLevelDomain = map[string]bool{
	"ac":  true,
	"co":  true,
	"org": true,
	"com": true,
	"net": true,
	"edu": true,
}

func host2Domain(host string) (domain string) {
	host, _ = splitHostPort(host)
	host = trimLastDot(host)
	lastDot := strings.LastIndex(host, ".")
	if lastDot == -1 {
		return host // simple host name, we should not hanlde this
	}
	// Find the 2nd last dot
	dot2ndLast := strings.LastIndex(host[:lastDot], ".")
	if dot2ndLast == -1 {
		return host
	}

	part := host[dot2ndLast+1 : lastDot]
	// If the 2nd last part of a domain name equals to a top level
	// domain, search for the 3rd part in the host name.
	// So domains like bbc.co.uk will not be recorded as co.uk
	if topLevelDomain[part] {
		dot3rdLast := strings.LastIndex(host[:dot2ndLast], ".")
		if dot3rdLast == -1 {
			return host
		}
		return host[dot3rdLast+1:]
	}
	return host[dot2ndLast+1:]
}

// TODO reload domain list when receives SIGUSR1
// one difficult here is that we may concurrently access maps which is not
// safe.
// Can we create a new domain set first, then change the reference of the original one?
// Domain set reference changing should be atomic.

func loadDomainSet() {
	blockedDs.loadDomainList(config.blockedFile)
	directDs.loadDomainList(config.directFile)
	alwaysBlockedDs.loadDomainList(config.alwaysBlockedFile)
	alwaysDirectDs.loadDomainList(config.alwaysDirectFile)
	chouDs.loadDomainList(config.chouFile)

	filterOutDs(chouDs)
	filterOutDs(alwaysDirectDs)
	filterOutDs(alwaysBlockedDs)
	filterOutBlockedDsInDirectDs()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/CleverbotIO/go-cleverbot.io"
	"github.com/PonyvilleFM/aura/pvfm/station"
	"github.com/bwmarrin/discordgo"
	"github.com/jzelinskie/geddit"
	"github.com/techniponi/doritobot/gitupdate"
)

func init() {
	rand.Seed(time.Now().Unix())
}

// randomRange gives a random whole integer between the given integers [min, max)
func randomRange(min, max int) int {
	return rand.Intn(max-min) + min
}

// Config is a json file containing the login data and any user-configurable values
type Config struct {
	Username            string `json:"username"`
	Password            string `json:"password"`
	Bot                 bool   `json:"Bot"`
	BotToken            string `json:"botToken"`
	HTTPEndpointEnabled bool   `json:"httpEndpointEnabled"`
	EchoCommandEnabled  bool   `json:"echoCommandEnabled"`
}

// DerpiResults is a struct to contain Derpibooru search results
type DerpiResults struct {
	Search []struct {
		ID               string        `json:"id"`
		CreatedAt        time.Time     `json:"created_at"`
		UpdatedAt        time.Time     `json:"updated_at"`
		DuplicateReports []interface{} `json:"duplicate_reports"`
		FirstSeenAt      time.Time     `json:"first_seen_at"`
		UploaderID       string        `json:"uploader_id"`
		Score            int           `json:"score"`
		CommentCount     int           `json:"comment_count"`
		Width            int           `json:"width"`
		Height           int           `json:"height"`
		FileName         string        `json:"file_name"`
		Description      string        `json:"description"`
		Uploader         string        `json:"uploader"`
		Image            string        `json:"image"`
		Upvotes          int           `json:"upvotes"`
		Downvotes        int           `json:"downvotes"`
		Faves            int           `json:"faves"`
		Tags             string        `json:"tags"`
		TagIds           []string      `json:"tag_ids"`
		AspectRatio      float64       `json:"aspect_ratio"`
		OriginalFormat   string        `json:"original_format"`
		MimeType         string        `json:"mime_type"`
		Sha512Hash       string        `json:"sha512_hash"`
		OrigSha512Hash   string        `json:"orig_sha512_hash"`
		SourceURL        string        `json:"source_url"`
		Representations  struct {
			ThumbTiny  string `json:"thumb_tiny"`
			ThumbSmall string `json:"thumb_small"`
			Thumb      string `json:"thumb"`
			Small      string `json:"small"`
			Medium     string `json:"medium"`
			Large      string `json:"large"`
			Tall       string `json:"tall"`
			Full       string `json:"full"`
		} `json:"representations"`
		IsRendered  bool `json:"is_rendered"`
		IsOptimized bool `json:"is_optimized"`
	} `json:"search"`
	Total        int           `json:"total"`
	Interactions []interface{} `json:"interactions"`
}

// These need to be global
var cb *cleverbot.Session
var botErr error
var currentID string
var userIDError error
var dg *discordgo.Session
var err error
var login Config

func main() {
	// Read the config.json
	loginFile, fileErr := ioutil.ReadFile("./config.json")
	if fileErr != nil {
		fmt.Println(fileErr.Error())
		os.Exit(1)
	}

	// Parse it
	jsonErr := json.Unmarshal(loginFile, &login)

	if jsonErr != nil {
		fmt.Println(jsonErr.Error())
		os.Exit(1)
	}

	// Log in to discord
	if login.Bot {
		dg, err = discordgo.New("Bot " + login.BotToken)
	} else {
		dg, err = discordgo.New(login.Username, login.Password)
	}

	if err != nil {
		fmt.Println("Error creating session: ", err)
		os.Exit(1)
	} else {
		fmt.Println("Successfully logged in.")
	}

	// Message received handler
	dg.AddHandler(messageCreate)

	err = dg.Open()
	if err != nil {
		fmt.Println("Error opening connection: ", err)
		os.Exit(1)
	}

	// Open cleverbot API
	// TODO: move this to login.json
	cb, botErr = cleverbot.New("JHowZe3ddT6Da0JU", "8NSK1vZVH1lRNMIcTbu4hU6kGEyIDxsW", "")
	if botErr != nil {
		log.Fatal(botErr)
	}

	// HTTP endpoint handler
	http.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
		if login.HTTPEndpointEnabled {
			id := r.URL.Query().Get("id")
			msg := r.URL.Query().Get("msg")
			if id != "" && msg != "" {
				dg.ChannelMessageSend(id, "Message from HTTP chat endpoint: \n`"+msg+"`")
				fmt.Fprintf(w, "You are sending: \n"+msg+"\nto channel ID: \n"+id)
			}
		}
	})

	http.HandleFunc("/repoupdate", func(w http.ResponseWriter, r *http.Request) {
		decoder := json.NewDecoder(r.Body)
		var update gitupdate.GitUpdate
		err := decoder.Decode(&update)
		if err != nil {
			log.Fatal(err)
		}
		defer r.Body.Close()
		updateString := "New commit pushed to <https://github.com/techniponi/doritobot> (" + update.HeadCommit.ID + "):\n" + update.HeadCommit.Author.Name + ": " + update.HeadCommit.Message + "\n" + update.HeadCommit.URL

		dg.ChannelMessageSend("298642620849324035", updateString)
	})

	dgUser, userIDError := dg.User("@me")
	currentID = dgUser.ID

	if userIDError != nil {
		log.Fatal(userIDError)
	}

	fmt.Println("Bot is now running.  Press CTRL-C to exit.")
	log.Fatal(http.ListenAndServe(":8080", nil))
	<-make(chan struct{})
	return
}

// messageCreate is the handler function for all incoming Discord messages
func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	var msg []string
	if m.Author.Username == "PonyChat" {
		// this is disabled since my implementation was SHIT
		//msg = strings.Split(m.Content[strings.Index(m.Content, ">")+4:len(m.Content)], " ")
	} else {
		msg = strings.Split(m.Content, " ")
	}

	if len(msg) > 0 {
		switch msg[0] {
		case "nazoupdate": // Reminds everyone that Nazo is adorable
			s.ChannelMessageDelete(m.ChannelID, m.ID)
			s.ChannelMessageSend(m.ChannelID, "This is your daily reminder that <@!165846085020024832> is adorable.")
			break
		case "echo": // Echoes the given text
			if login.EchoCommandEnabled {
				s.ChannelMessageDelete(m.ChannelID, m.ID)
				/*if m.Author.Username == "PonyChat" {
					s.ChannelMessageSend(m.ChannelID, "`ds:` "+m.Content[15:len(m.Content)])
				} else {*/
				s.ChannelMessageSend(m.ChannelID, "`echo:`\n"+m.Content[11:len(m.Content)])
				//}
			}
			break
		case "cb": // Communicates with Cleverbot.
			response, botErr2 := cb.Ask(m.Content[2:len(m.Content)])
			if botErr2 != nil {
				log.Fatal(botErr2)
			}
			s.ChannelMessageSend(m.ChannelID, response)
			break
		case "db": // Grabs an image from Derpibooru results with a given list of tags - always a safe image!
			if len(msg) < 2 {
				s.ChannelMessageSend(m.ChannelID, "Error: not enough arguments")
				break
			}
			derpiTags := strings.Replace(m.Content[3:len(m.Content)], " ", "+", -1)
			resp, derpiHTTPErr := http.Get("https://derpibooru.org/search.json?q=safe," + derpiTags)
			if derpiHTTPErr != nil {
				s.ChannelMessageSend(m.ChannelID, "HTTP error, check console")
				log.Fatal(derpiHTTPErr)
				break
			}
			defer resp.Body.Close()
			derpiBody, derpiErr := ioutil.ReadAll(resp.Body)
			var derpiErr2 error
			results := DerpiResults{}
			derpiErr2 = json.Unmarshal(derpiBody, &results)
			if derpiErr != nil {
				s.ChannelMessageSend(m.ChannelID, "Error reading HTTP response")
				break
			}
			if derpiErr2 != nil {
				s.ChannelMessageSend(m.ChannelID, "Error reading derpibooru API response")
				break
			}
			if len(results.Search) > 0 {
				s.ChannelMessageSend(m.ChannelID, "http:"+results.Search[randomRange(0, len(results.Search))].Image)
			} else {
				s.ChannelMessageSend(m.ChannelID, "Error: no results")
			}
			break
		case "h": // h
			if m.Author.ID != currentID {
				s.ChannelMessageSend(m.ChannelID, "h")
			}
			break
		case "pvfmservers": // Gives a list of all available PVFM streams (direct links)
			currentMeta, metaErr := station.GetStats()
			if metaErr != nil {
				s.ChannelMessageSend(m.ChannelID, "Error receiving pvfm metadata")
				break
			}

			outputString := "**PVFM Servers:**\n"

			for _, element := range currentMeta.Icestats.Source {
				outputString += ":musical_note: " + element.ServerDescription + ":\n`" + strings.Replace(element.Listenurl, "aerial", "dj.bronyradio.com", -1) + "`\n"
			}

			outputString += "\n:cd: DJ Recordings:\n`http://darkling.darkwizards.com/wang/BronyRadio/?M=D`"

			s.ChannelMessageSend(m.ChannelID, outputString)

			break
		case "techgore":
			listOptions := geddit.ListingOptions{
				Limit: 50,
			}
			reddit := geddit.NewSession("discordbot")
			results, err := reddit.SubredditSubmissions("techsupportgore", geddit.NewSubmissions, listOptions)
			if err != nil {
				break
			}
			s.ChannelMessageSend(m.ChannelID, results[randomRange(0, len(results))].URL)
			break
		case "snuggle", "cuddle", "hug", "kiss", "boop", "glomp", "nuzzle":
			if len(msg) < 2 {
				s.ChannelMessageSend(m.ChannelID, "Who - Delta, Twisty, or Jac?")
				break
			}
			names := map[string]string{
				"awal":    "Twisty",
				"twisty":  "Twisty",
				"delta":   "Delta",
				"dorito":  "Delta",
				"techni":  "Delta",
				"nazo":    "Jac",
				"nuzzles": "Jac",
				"jac":     "Jac",
				"thorax":  "Thorax",
				"shining": "Shiny",
				"shiny":   "Shiny",
				"quartz":  "Quartz",
				"dyed":    "Quartz",
			}
			possibleResponses := []string{
				"\"ees\" softly.",
				"snuggles back.",
				"flops over.",
				"blushes profusely.",
				"twitches his ears and smiles.",
				"smiles lovingly.",
				"boops you back!",
				"glomps you!",
				"is happy.",
				"jumps with joy!",
				"wasn't expecting that! :heart:",
				"loves you.",
			}
			characterSpecifics := map[string][]string{
				"Thorax": {"vibrates his wings in excitement.", "is cheered up from your kindness!"},
				"Shiny":  {"wonders if Cadance is okay with this.", "thinks you would be a great addition to the Sparkle family."},
				"Delta":  {"gets a wingboner.", "vibrates."},
				"Jac":    {"dies of cuteness overload.", "passes out from an extreme overdose of gay.", "can't hold all these husbandos."},
				"Twisty": {"invites you to his next gig.", "needed that! :heart:"},
				"Quartz": {"runs away.", "did not like that."},
			}
			if names[msg[1]] == "" {
				s.ChannelMessageSend(m.ChannelID, "I'm afraid I don't know who that is. :c")
				break
			}
			finalMessage := "error" // set to error as default in case of derpage
			if randomRange(0, 10) == 7 {
				finalMessage = characterSpecifics[names[msg[1]]][randomRange(0, len(characterSpecifics[names[msg[1]]]))] // this line is a fucking mess
			} else {
				finalMessage = possibleResponses[randomRange(0, len(possibleResponses))]
			}
			s.ChannelMessageSend(m.ChannelID, names[msg[1]]+" "+finalMessage)

			break
		case "gay":
			possibleResponses := []string{
				"floof.zone/img/gay.jpg",
				"floof.zone/img/nazoblep.png",
				"floof.zone/img/awalblep.png",
				"floof.zone/img/floofbat-capall.png",
			}
			s.ChannelMessageSend(m.ChannelID, "http://"+possibleResponses[randomRange(0, len(possibleResponses))])
			break
		}
	}
}
<commit_msg>making my character snuggle better ;)<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/CleverbotIO/go-cleverbot.io"
	"github.com/PonyvilleFM/aura/pvfm/station"
	"github.com/bwmarrin/discordgo"
	"github.com/jzelinskie/geddit"
	"github.com/techniponi/doritobot/gitupdate"
)

func init() {
	rand.Seed(time.Now().Unix())
}

// randomRange gives a random whole integer between the given integers [min, max)
func randomRange(min, max int) int {
	return rand.Intn(max-min) + min
}

// Config is a json file containing the login data and any user-configurable values
type Config struct {
	Username            string `json:"username"`
	Password            string `json:"password"`
	Bot                 bool   `json:"Bot"`
	BotToken            string `json:"botToken"`
	HTTPEndpointEnabled bool   `json:"httpEndpointEnabled"`
	EchoCommandEnabled  bool   `json:"echoCommandEnabled"`
}

// DerpiResults is a struct to contain Derpibooru search results
type DerpiResults struct {
	Search []struct {
		ID               string        `json:"id"`
		CreatedAt        time.Time     `json:"created_at"`
		UpdatedAt        time.Time     `json:"updated_at"`
		DuplicateReports []interface{} `json:"duplicate_reports"`
		FirstSeenAt      time.Time     `json:"first_seen_at"`
		UploaderID       string        `json:"uploader_id"`
		Score            int           `json:"score"`
		CommentCount     int           `json:"comment_count"`
		Width            int           `json:"width"`
		Height           int           `json:"height"`
		FileName         string        `json:"file_name"`
		Description      string        `json:"description"`
		Uploader         string        `json:"uploader"`
		Image            string        `json:"image"`
		Upvotes          int           `json:"upvotes"`
		Downvotes        int           `json:"downvotes"`
		Faves            int           `json:"faves"`
		Tags             string        `json:"tags"`
		TagIds           []string      `json:"tag_ids"`
		AspectRatio      float64       `json:"aspect_ratio"`
		OriginalFormat   string        `json:"original_format"`
		MimeType         string        `json:"mime_type"`
		Sha512Hash       string        `json:"sha512_hash"`
		OrigSha512Hash   string        `json:"orig_sha512_hash"`
		SourceURL        string        `json:"source_url"`
		Representations  struct {
			ThumbTiny  string `json:"thumb_tiny"`
			ThumbSmall string `json:"thumb_small"`
			Thumb      string `json:"thumb"`
			Small      string `json:"small"`
			Medium     string `json:"medium"`
			Large      string `json:"large"`
			Tall       string `json:"tall"`
			Full       string `json:"full"`
		} `json:"representations"`
		IsRendered  bool `json:"is_rendered"`
		IsOptimized bool `json:"is_optimized"`
	} `json:"search"`
	Total        int           `json:"total"`
	Interactions []interface{} `json:"interactions"`
}

// These need to be global
var cb *cleverbot.Session
var botErr error
var currentID string
var userIDError error
var dg *discordgo.Session
var err error
var login Config

func main() {
	// Read the config.json
	loginFile, fileErr := ioutil.ReadFile("./config.json")
	if fileErr != nil {
		fmt.Println(fileErr.Error())
		os.Exit(1)
	}

	// Parse it
	jsonErr := json.Unmarshal(loginFile, &login)

	if jsonErr != nil {
		fmt.Println(jsonErr.Error())
		os.Exit(1)
	}

	// Log in to discord
	if login.Bot {
		dg, err = discordgo.New("Bot " + login.BotToken)
	} else {
		dg, err = discordgo.New(login.Username, login.Password)
	}

	if err != nil {
		fmt.Println("Error creating session: ", err)
		os.Exit(1)
	} else {
		fmt.Println("Successfully logged in.")
	}

	// Message received handler
	dg.AddHandler(messageCreate)

	err = dg.Open()
	if err != nil {
		fmt.Println("Error opening connection: ", err)
		os.Exit(1)
	}

	// Open cleverbot API
	// TODO: move this to login.json
	cb, botErr = cleverbot.New("JHowZe3ddT6Da0JU", "8NSK1vZVH1lRNMIcTbu4hU6kGEyIDxsW", "")
	if botErr != nil {
		log.Fatal(botErr)
	}

	// HTTP endpoint handler
	http.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
		if login.HTTPEndpointEnabled {
			id := r.URL.Query().Get("id")
			msg := r.URL.Query().Get("msg")
			if id != "" && msg != "" {
				dg.ChannelMessageSend(id, "Message from HTTP chat endpoint: \n`"+msg+"`")
				fmt.Fprintf(w, "You are sending: \n"+msg+"\nto channel ID: \n"+id)
			}
		}
	})

	http.HandleFunc("/repoupdate", func(w http.ResponseWriter, r *http.Request) {
		decoder := json.NewDecoder(r.Body)
		var update gitupdate.GitUpdate
		err := decoder.Decode(&update)
		if err != nil {
			log.Fatal(err)
		}
		defer r.Body.Close()
		updateString := "New commit pushed to <https://github.com/techniponi/doritobot> (" + update.HeadCommit.ID + "):\n" + update.HeadCommit.Author.Name + ": " + update.HeadCommit.Message + "\n" + update.HeadCommit.URL

		dg.ChannelMessageSend("298642620849324035", updateString)
	})

	dgUser, userIDError := dg.User("@me")
	currentID = dgUser.ID

	if userIDError != nil {
		log.Fatal(userIDError)
	}

	fmt.Println("Bot is now running.  Press CTRL-C to exit.")
	log.Fatal(http.ListenAndServe(":8080", nil))
	<-make(chan struct{})
	return
}

// messageCreate is the handler function for all incoming Discord messages
func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	var msg []string
	if m.Author.Username == "PonyChat" {
		// this is disabled since my implementation was SHIT
		//msg = strings.Split(m.Content[strings.Index(m.Content, ">")+4:len(m.Content)], " ")
	} else {
		msg = strings.Split(m.Content, " ")
	}

	if len(msg) > 0 {
		switch msg[0] {
		case "nazoupdate": // Reminds everyone that Nazo is adorable
			s.ChannelMessageDelete(m.ChannelID, m.ID)
			s.ChannelMessageSend(m.ChannelID, "This is your daily reminder that <@!165846085020024832> is adorable.")
			break
		case "echo": // Echoes the given text
			if login.EchoCommandEnabled {
				s.ChannelMessageDelete(m.ChannelID, m.ID)
				/*if m.Author.Username == "PonyChat" {
					s.ChannelMessageSend(m.ChannelID, "`ds:` "+m.Content[15:len(m.Content)])
				} else {*/
				s.ChannelMessageSend(m.ChannelID, "`echo:`\n"+m.Content[11:len(m.Content)])
				//}
			}
			break
		case "cb": // Communicates with Cleverbot.
			response, botErr2 := cb.Ask(m.Content[2:len(m.Content)])
			if botErr2 != nil {
				log.Fatal(botErr2)
			}
			s.ChannelMessageSend(m.ChannelID, response)
			break
		case "db": // Grabs an image from Derpibooru results with a given list of tags - always a safe image!
			if len(msg) < 2 {
				s.ChannelMessageSend(m.ChannelID, "Error: not enough arguments")
				break
			}
			derpiTags := strings.Replace(m.Content[3:len(m.Content)], " ", "+", -1)
			resp, derpiHTTPErr := http.Get("https://derpibooru.org/search.json?q=safe," + derpiTags)
			if derpiHTTPErr != nil {
				s.ChannelMessageSend(m.ChannelID, "HTTP error, check console")
				log.Fatal(derpiHTTPErr)
				break
			}
			defer resp.Body.Close()
			derpiBody, derpiErr := ioutil.ReadAll(resp.Body)
			var derpiErr2 error
			results := DerpiResults{}
			derpiErr2 = json.Unmarshal(derpiBody, &results)
			if derpiErr != nil {
				s.ChannelMessageSend(m.ChannelID, "Error reading HTTP response")
				break
			}
			if derpiErr2 != nil {
				s.ChannelMessageSend(m.ChannelID, "Error reading derpibooru API response")
				break
			}
			if len(results.Search) > 0 {
				s.ChannelMessageSend(m.ChannelID, "http:"+results.Search[randomRange(0, len(results.Search))].Image)
			} else {
				s.ChannelMessageSend(m.ChannelID, "Error: no results")
			}
			break
		case "h": // h
			if m.Author.ID != currentID {
				s.ChannelMessageSend(m.ChannelID, "h")
			}
			break
		case "pvfmservers": // Gives a list of all available PVFM streams (direct links)
			currentMeta, metaErr := station.GetStats()
			if metaErr != nil {
				s.ChannelMessageSend(m.ChannelID, "Error receiving pvfm metadata")
				break
			}

			outputString := "**PVFM Servers:**\n"

			for _, element := range currentMeta.Icestats.Source {
				outputString += ":musical_note: " + element.ServerDescription + ":\n`" + strings.Replace(element.Listenurl, "aerial", "dj.bronyradio.com", -1) + "`\n"
			}

			outputString += "\n:cd: DJ Recordings:\n`http://darkling.darkwizards.com/wang/BronyRadio/?M=D`"

			s.ChannelMessageSend(m.ChannelID, outputString)

			break
		case "techgore":
			listOptions := geddit.ListingOptions{
				Limit: 50,
			}
			reddit := geddit.NewSession("discordbot")
			results, err := reddit.SubredditSubmissions("techsupportgore", geddit.NewSubmissions, listOptions)
			if err != nil {
				break
			}
			s.ChannelMessageSend(m.ChannelID, results[randomRange(0, len(results))].URL)
			break
		case "snuggle", "cuddle", "hug", "kiss", "boop", "glomp", "nuzzle":
			if len(msg) < 2 {
				s.ChannelMessageSend(m.ChannelID, "Who - Delta, Twisty, or Jac?")
				break
			}
			names := map[string]string{
				"awal":    "Twisty",
				"twisty":  "Twisty",
				"delta":   "Delta",
				"dorito":  "Delta",
				"techni":  "Delta",
				"nazo":    "Jac",
				"nuzzles": "Jac",
				"jac":     "Jac",
				"thorax":  "Thorax",
				"shining": "Shiny",
				"shiny":   "Shiny",
				"quartz":  "Quartz",
				"dyed":    "Quartz",
			}
			possibleResponses := []string{
				"\"ees\" softly.",
				"snuggles back.",
				"flops over.",
				"blushes profusely.",
				"twitches his ears and smiles.",
				"smiles lovingly.",
				"boops you back!",
				"glomps you!",
				"is happy.",
				"jumps with joy!",
				"wasn't expecting that! :heart:",
				"loves you.",
			}
			characterSpecifics := map[string][]string{
				"Thorax": {"vibrates his wings in excitement.", "is cheered up from your kindness!"},
				"Shiny":  {"wonders if Cadance is okay with this.", "thinks you would be a great addition to the Sparkle family."},
				"Delta":  {"gets a wingboner.", "vibrates."},
				"Jac":    {"dies of cuteness overload.", "passes out from an extreme overdose of gay.", "can't hold all these husbandos."},
				"Twisty": {"invites you to his next gig.", "needed that! :heart:"},
				"Quartz": {"runs away.", "did not like that."},
			}
			if names[msg[1]] == "" {
				s.ChannelMessageSend(m.ChannelID, "I'm afraid I don't know who that is. :c")
				break
			}
			finalMessage := "error" // set to error as default in case of derpage

			if names[msg[1]] == "Quartz" {
				finalMessage = characterSpecifics[names[msg[1]]][randomRange(0, len(characterSpecifics[names[msg[1]]]))]
			} else if randomRange(0, 10) == 7 {
				finalMessage = characterSpecifics[names[msg[1]]][randomRange(0, len(characterSpecifics[names[msg[1]]]))] // this line is a fucking mess
			} else {
				finalMessage = possibleResponses[randomRange(0, len(possibleResponses))]
			}

			s.ChannelMessageSend(m.ChannelID, names[msg[1]]+" "+finalMessage)

			break
		case "gay":
			possibleResponses := []string{
				"floof.zone/img/gay.jpg",
				"floof.zone/img/nazoblep.png",
				"floof.zone/img/awalblep.png",
				"floof.zone/img/floofbat-capall.png",
			}
			s.ChannelMessageSend(m.ChannelID, "http://"+possibleResponses[randomRange(0, len(possibleResponses))])
			break
		}
	}
}
<|endoftext|>