<commit_before>package utp

/*
#include "utp.h"
#include <stdbool.h>

struct utp_process_udp_args {
	const byte *buf;
	size_t len;
	const struct sockaddr *sa;
	socklen_t sal;
};

void process_received_messages(utp_context *ctx, struct utp_process_udp_args *args, size_t argslen)
{
	bool gotUtp = false;
	size_t i;
	for (i = 0; i < argslen; i++) {
		struct utp_process_udp_args *a = &args[i];
		//if (!a->len) continue;
		if (utp_process_udp(ctx, a->buf, a->len, a->sa, a->sal)) {
			gotUtp = true;
		}
	}
	if (gotUtp) {
		utp_issue_deferred_acks(ctx);
		utp_check_timeouts(ctx);
	}
}
*/
import "C"
import (
	"context"
	"errors"
	"net"
	"time"

	"github.com/anacrolix/missinggo"
	"github.com/anacrolix/missinggo/inproc"
	"github.com/anacrolix/mmsg"
)

type Socket struct {
	pc               net.PacketConn
	ctx              *C.utp_context
	backlog          chan *Conn
	closed           bool
	conns            map[*C.utp_socket]*Conn
	nonUtpReads      chan packet
	writeDeadline    time.Time
	readDeadline     time.Time
	firewallCallback FirewallCallback
	// Whether the next accept is to be blocked.
	block bool
}

type FirewallCallback func(net.Addr) bool

var (
	_               net.PacketConn = (*Socket)(nil)
	_               net.Listener   = (*Socket)(nil)
	errSocketClosed                = errors.New("Socket closed")
)

type packet struct {
	b    []byte
	from net.Addr
}

func listenPacket(network, addr string) (pc net.PacketConn, err error) {
	if network == "inproc" {
		return inproc.ListenPacket(network, addr)
	}
	return net.ListenPacket(network, addr)
}

func NewSocket(network, addr string) (*Socket, error) {
	pc, err := listenPacket(network, addr)
	if err != nil {
		return nil, err
	}
	mu.Lock()
	defer mu.Unlock()
	ctx := C.utp_init(2)
	if ctx == nil {
		panic(ctx)
	}
	ctx.setCallbacks()
	if utpLogging {
		ctx.setOption(C.UTP_LOG_NORMAL, 1)
		ctx.setOption(C.UTP_LOG_MTU, 1)
		ctx.setOption(C.UTP_LOG_DEBUG, 1)
	}
	s := &Socket{
		pc:          pc,
		ctx:         ctx,
		backlog:     make(chan *Conn, 5),
		conns:       make(map[*C.utp_socket]*Conn),
		nonUtpReads: make(chan packet, 100),
	}
	libContextToSocket[ctx] = s
	go s.timeoutChecker()
	go s.packetReader()
	return s, nil
}

func (s *Socket) onLibSocketDestroyed(ls *C.utp_socket) {
	delete(s.conns, ls)
}

func (s *Socket) newConn(us *C.utp_socket) *Conn {
	c := &Conn{
		s:         s,
		us:        us,
		localAddr: s.pc.LocalAddr(),
	}
	c.cond.L = &mu
	s.conns[us] = c
	c.writeDeadlineTimer = time.AfterFunc(-1, c.cond.Broadcast)
	c.readDeadlineTimer = time.AfterFunc(-1, c.cond.Broadcast)
	return c
}

const maxNumBuffers = 16

func (s *Socket) packetReader() {
	mc := mmsg.NewConn(s.pc)
	// Increasing the messages increases the memory use, but also means we can
	// reduces utp_issue_deferred_acks and syscalls which should improve
	// efficiency. On the flip side, not all OSs implement batched reads.
	ms := make([]mmsg.Message, func() int {
		if mc.Err() == nil {
			return maxNumBuffers
		} else {
			return 1
		}
	}())
	for i := range ms {
		// The IPv4 UDP limit is allegedly about 64 KiB, and this message has
		// been seen on receiving on Windows with just 0x1000: wsarecvfrom: A
		// message sent on a datagram socket was larger than the internal
		// message buffer or some other network limit, or the buffer used to
		// receive a datagram into was smaller than the datagram itself.
		ms[i].Buffers = [][]byte{make([]byte, 0x10000)}
	}
	// Some crap OSs like Windoze will raise errors in Reads that don't
	// actually mean we should stop.
	consecutiveErrors := 0
	for {
		// In C, all the reads are processed and when it threatens to block,
		// we're supposed to call utp_issue_deferred_acks.
		n, err := mc.RecvMsgs(ms)
		if n == 1 {
			singleMsgRecvs.Add(1)
		}
		if n > 1 {
			multiMsgRecvs.Add(1)
		}
		if err != nil {
			mu.Lock()
			closed := s.closed
			mu.Unlock()
			if closed {
				// We don't care.
				return
			}
			// See https://github.com/anacrolix/torrent/issues/83. If we get
			// an endless stream of errors (such as the PacketConn being
			// Closed outside of our control, this work around may need to be
			// reconsidered.
			Logger.Printf("ignoring socket read error: %s", err)
			consecutiveErrors++
			if consecutiveErrors >= 100 {
				Logger.Print("too many consecutive errors, closing socket")
				s.Close()
				return
			}
			continue
		}
		consecutiveErrors = 0
		expMap.Add("successful mmsg receive calls", 1)
		expMap.Add("received messages", int64(n))
		s.processReceivedMessages(ms[:n])
	}
}

func (s *Socket) processReceivedMessages(ms []mmsg.Message) {
	mu.Lock()
	defer mu.Unlock()
	if s.closed {
		return
	}
	if processPacketsInC {
		var args [maxNumBuffers]C.struct_utp_process_udp_args
		for i, m := range ms {
			a := &args[i]
			a.buf = (*C.byte)(&m.Buffers[0][0])
			a.len = C.size_t(m.N)
			a.sa, a.sal = netAddrToLibSockaddr(m.Addr)
		}
		C.process_received_messages(s.ctx, &args[0], C.size_t(len(ms)))
	} else {
		gotUtp := false
		for _, m := range ms {
			gotUtp = s.processReceivedMessage(m.Buffers[0][:m.N], m.Addr) || gotUtp
		}
		if gotUtp && !s.closed {
			s.afterReceivingUtpMessages()
		}
	}
}

func (s *Socket) afterReceivingUtpMessages() {
	C.utp_issue_deferred_acks(s.ctx)
	// TODO: When is this done in C?
	C.utp_check_timeouts(s.ctx)
}

func (s *Socket) processReceivedMessage(b []byte, addr net.Addr) (utp bool) {
	if s.utpProcessUdp(b, addr) {
		socketUtpPacketsReceived.Add(1)
		return true
	} else {
		s.onReadNonUtp(b, addr)
		return false
	}
}

// Process packet batches entirely from C, reducing CGO overhead. Currently
// requires GODEBUG=cgocheck=0.
const processPacketsInC = false

// Wraps libutp's utp_process_udp, returning relevant information.
func (s *Socket) utpProcessUdp(b []byte, addr net.Addr) (utp bool) {
	sa, sal := netAddrToLibSockaddr(addr)
	if len(b) == 0 {
		// The implementation of utp_process_udp rejects null buffers, and
		// anything smaller than the UTP header size. It's also prone to
		// assert on those, which we don't want to trigger.
		return false
	}
	if missinggo.AddrPort(addr) == 0 {
		return false
	}
	mu.Unlock()
	block := func() bool {
		if s.firewallCallback == nil {
			return false
		}
		return s.firewallCallback(addr)
	}()
	mu.Lock()
	s.block = block
	if s.closed {
		return false
	}
	ret := C.utp_process_udp(s.ctx, (*C.byte)(&b[0]), C.size_t(len(b)), sa, sal)
	switch ret {
	case 1:
		return true
	case 0:
		return false
	default:
		panic(ret)
	}
}

func (s *Socket) timeoutChecker() {
	for {
		mu.Lock()
		if s.closed {
			mu.Unlock()
			return
		}
		// C.utp_issue_deferred_acks(s.ctx)
		C.utp_check_timeouts(s.ctx)
		mu.Unlock()
		time.Sleep(500 * time.Millisecond)
	}
}

func (s *Socket) Close() error {
	mu.Lock()
	defer mu.Unlock()
	return s.closeLocked()
}

func (s *Socket) closeLocked() error {
	if s.closed {
		return nil
	}
	// Calling this deletes the pointer. It must not be referred to after
	// this.
	C.utp_destroy(s.ctx)
	s.ctx = nil
	s.pc.Close()
	close(s.backlog)
	close(s.nonUtpReads)
	s.closed = true
	return nil
}

func (s *Socket) Addr() net.Addr {
	return s.pc.LocalAddr()
}

func (s *Socket) LocalAddr() net.Addr {
	return s.pc.LocalAddr()
}

func (s *Socket) Accept() (net.Conn, error) {
	nc, ok := <-s.backlog
	if !ok {
		return nil, errors.New("closed")
	}
	return nc, nil
}

func (s *Socket) Dial(addr string) (net.Conn, error) {
	return s.DialTimeout(addr, 0)
}

func (s *Socket) DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
	ctx := context.Background()
	if timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel()
	}
	return s.DialContext(ctx, "", addr)
}

func (s *Socket) resolveAddr(network, addr string) (net.Addr, error) {
	if network == "" {
		network = s.Addr().Network()
	}
	return resolveAddr(network, addr)
}

func resolveAddr(network, addr string) (net.Addr, error) {
	switch network {
	case "inproc":
		return inproc.ResolveAddr(network, addr)
	default:
		return net.ResolveUDPAddr(network, addr)
	}
}

// Passing an empty network will use the network of the Socket's listener.
func (s *Socket) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
	c, err := s.NewConn()
	if err != nil {
		return nil, err
	}
	err = c.Connect(ctx, network, addr)
	if err != nil {
		c.Close()
		return nil, err
	}
	return c, nil
}

func (s *Socket) NewConn() (*Conn, error) {
	mu.Lock()
	defer mu.Unlock()
	if s.closed {
		return nil, errors.New("socket closed")
	}
	return s.newConn(C.utp_create_socket(s.ctx)), nil
}

func (s *Socket) pushBacklog(c *Conn) {
	select {
	case s.backlog <- c:
	default:
		c.close()
	}
}

func (s *Socket) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	p, ok := <-s.nonUtpReads
	if !ok {
		err = errors.New("closed")
		return
	}
	n = copy(b, p.b)
	addr = p.from
	return
}

func (s *Socket) onReadNonUtp(b []byte, from net.Addr) {
	if s.closed {
		return
	}
	socketNonUtpPacketsReceived.Add(1)
	select {
	case s.nonUtpReads <- packet{append([]byte(nil), b...), from}:
	default:
		// log.Printf("dropped non utp packet: no room in buffer")
		nonUtpPacketsDropped.Add(1)
	}
}

func (s *Socket) SetReadDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) SetWriteDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) SetDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) WriteTo(b []byte, addr net.Addr) (int, error) {
	return s.pc.WriteTo(b, addr)
}

func (s *Socket) ReadBufferLen() int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_get_option(s.ctx, C.UTP_RCVBUF))
}

func (s *Socket) WriteBufferLen() int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_get_option(s.ctx, C.UTP_SNDBUF))
}

func (s *Socket) SetWriteBufferLen(len int) {
	mu.Lock()
	defer mu.Unlock()
	i := C.utp_context_set_option(s.ctx, C.UTP_SNDBUF, C.int(len))
	if i != 0 {
		panic(i)
	}
}

func (s *Socket) SetOption(opt Option, val int) int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_set_option(s.ctx, opt, C.int(val)))
}

func (s *Socket) SetFirewallCallback(f FirewallCallback) {
	mu.Lock()
	s.firewallCallback = f
	mu.Unlock()
}
<commit_msg>Minimize time spent holding lock in NewSocket<commit_after>package utp

/*
#include "utp.h"
#include <stdbool.h>

struct utp_process_udp_args {
	const byte *buf;
	size_t len;
	const struct sockaddr *sa;
	socklen_t sal;
};

void process_received_messages(utp_context *ctx, struct utp_process_udp_args *args, size_t argslen)
{
	bool gotUtp = false;
	size_t i;
	for (i = 0; i < argslen; i++) {
		struct utp_process_udp_args *a = &args[i];
		//if (!a->len) continue;
		if (utp_process_udp(ctx, a->buf, a->len, a->sa, a->sal)) {
			gotUtp = true;
		}
	}
	if (gotUtp) {
		utp_issue_deferred_acks(ctx);
		utp_check_timeouts(ctx);
	}
}
*/
import "C"
import (
	"context"
	"errors"
	"net"
	"time"

	"github.com/anacrolix/missinggo"
	"github.com/anacrolix/missinggo/inproc"
	"github.com/anacrolix/mmsg"
)

type Socket struct {
	pc               net.PacketConn
	ctx              *C.utp_context
	backlog          chan *Conn
	closed           bool
	conns            map[*C.utp_socket]*Conn
	nonUtpReads      chan packet
	writeDeadline    time.Time
	readDeadline     time.Time
	firewallCallback FirewallCallback
	// Whether the next accept is to be blocked.
	block bool
}

type FirewallCallback func(net.Addr) bool

var (
	_               net.PacketConn = (*Socket)(nil)
	_               net.Listener   = (*Socket)(nil)
	errSocketClosed                = errors.New("Socket closed")
)

type packet struct {
	b    []byte
	from net.Addr
}

func listenPacket(network, addr string) (pc net.PacketConn, err error) {
	if network == "inproc" {
		return inproc.ListenPacket(network, addr)
	}
	return net.ListenPacket(network, addr)
}

func NewSocket(network, addr string) (*Socket, error) {
	pc, err := listenPacket(network, addr)
	if err != nil {
		return nil, err
	}
	s := &Socket{
		pc:          pc,
		backlog:     make(chan *Conn, 5),
		conns:       make(map[*C.utp_socket]*Conn),
		nonUtpReads: make(chan packet, 100),
	}
	func() {
		mu.Lock()
		defer mu.Unlock()
		ctx := C.utp_init(2)
		if ctx == nil {
			panic(ctx)
		}
		s.ctx = ctx
		ctx.setCallbacks()
		if utpLogging {
			ctx.setOption(C.UTP_LOG_NORMAL, 1)
			ctx.setOption(C.UTP_LOG_MTU, 1)
			ctx.setOption(C.UTP_LOG_DEBUG, 1)
		}
		libContextToSocket[ctx] = s
	}()
	go s.timeoutChecker()
	go s.packetReader()
	return s, nil
}

func (s *Socket) onLibSocketDestroyed(ls *C.utp_socket) {
	delete(s.conns, ls)
}

func (s *Socket) newConn(us *C.utp_socket) *Conn {
	c := &Conn{
		s:         s,
		us:        us,
		localAddr: s.pc.LocalAddr(),
	}
	c.cond.L = &mu
	s.conns[us] = c
	c.writeDeadlineTimer = time.AfterFunc(-1, c.cond.Broadcast)
	c.readDeadlineTimer = time.AfterFunc(-1, c.cond.Broadcast)
	return c
}

const maxNumBuffers = 16

func (s *Socket) packetReader() {
	mc := mmsg.NewConn(s.pc)
	// Increasing the messages increases the memory use, but also means we can
	// reduces utp_issue_deferred_acks and syscalls which should improve
	// efficiency. On the flip side, not all OSs implement batched reads.
	ms := make([]mmsg.Message, func() int {
		if mc.Err() == nil {
			return maxNumBuffers
		} else {
			return 1
		}
	}())
	for i := range ms {
		// The IPv4 UDP limit is allegedly about 64 KiB, and this message has
		// been seen on receiving on Windows with just 0x1000: wsarecvfrom: A
		// message sent on a datagram socket was larger than the internal
		// message buffer or some other network limit, or the buffer used to
		// receive a datagram into was smaller than the datagram itself.
		ms[i].Buffers = [][]byte{make([]byte, 0x10000)}
	}
	// Some crap OSs like Windoze will raise errors in Reads that don't
	// actually mean we should stop.
	consecutiveErrors := 0
	for {
		// In C, all the reads are processed and when it threatens to block,
		// we're supposed to call utp_issue_deferred_acks.
		n, err := mc.RecvMsgs(ms)
		if n == 1 {
			singleMsgRecvs.Add(1)
		}
		if n > 1 {
			multiMsgRecvs.Add(1)
		}
		if err != nil {
			mu.Lock()
			closed := s.closed
			mu.Unlock()
			if closed {
				// We don't care.
				return
			}
			// See https://github.com/anacrolix/torrent/issues/83. If we get
			// an endless stream of errors (such as the PacketConn being
			// Closed outside of our control, this work around may need to be
			// reconsidered.
			Logger.Printf("ignoring socket read error: %s", err)
			consecutiveErrors++
			if consecutiveErrors >= 100 {
				Logger.Print("too many consecutive errors, closing socket")
				s.Close()
				return
			}
			continue
		}
		consecutiveErrors = 0
		expMap.Add("successful mmsg receive calls", 1)
		expMap.Add("received messages", int64(n))
		s.processReceivedMessages(ms[:n])
	}
}

func (s *Socket) processReceivedMessages(ms []mmsg.Message) {
	mu.Lock()
	defer mu.Unlock()
	if s.closed {
		return
	}
	if processPacketsInC {
		var args [maxNumBuffers]C.struct_utp_process_udp_args
		for i, m := range ms {
			a := &args[i]
			a.buf = (*C.byte)(&m.Buffers[0][0])
			a.len = C.size_t(m.N)
			a.sa, a.sal = netAddrToLibSockaddr(m.Addr)
		}
		C.process_received_messages(s.ctx, &args[0], C.size_t(len(ms)))
	} else {
		gotUtp := false
		for _, m := range ms {
			gotUtp = s.processReceivedMessage(m.Buffers[0][:m.N], m.Addr) || gotUtp
		}
		if gotUtp && !s.closed {
			s.afterReceivingUtpMessages()
		}
	}
}

func (s *Socket) afterReceivingUtpMessages() {
	C.utp_issue_deferred_acks(s.ctx)
	// TODO: When is this done in C?
	C.utp_check_timeouts(s.ctx)
}

func (s *Socket) processReceivedMessage(b []byte, addr net.Addr) (utp bool) {
	if s.utpProcessUdp(b, addr) {
		socketUtpPacketsReceived.Add(1)
		return true
	} else {
		s.onReadNonUtp(b, addr)
		return false
	}
}

// Process packet batches entirely from C, reducing CGO overhead. Currently
// requires GODEBUG=cgocheck=0.
const processPacketsInC = false

// Wraps libutp's utp_process_udp, returning relevant information.
func (s *Socket) utpProcessUdp(b []byte, addr net.Addr) (utp bool) {
	sa, sal := netAddrToLibSockaddr(addr)
	if len(b) == 0 {
		// The implementation of utp_process_udp rejects null buffers, and
		// anything smaller than the UTP header size. It's also prone to
		// assert on those, which we don't want to trigger.
		return false
	}
	if missinggo.AddrPort(addr) == 0 {
		return false
	}
	mu.Unlock()
	block := func() bool {
		if s.firewallCallback == nil {
			return false
		}
		return s.firewallCallback(addr)
	}()
	mu.Lock()
	s.block = block
	if s.closed {
		return false
	}
	ret := C.utp_process_udp(s.ctx, (*C.byte)(&b[0]), C.size_t(len(b)), sa, sal)
	switch ret {
	case 1:
		return true
	case 0:
		return false
	default:
		panic(ret)
	}
}

func (s *Socket) timeoutChecker() {
	for {
		mu.Lock()
		if s.closed {
			mu.Unlock()
			return
		}
		// C.utp_issue_deferred_acks(s.ctx)
		C.utp_check_timeouts(s.ctx)
		mu.Unlock()
		time.Sleep(500 * time.Millisecond)
	}
}

func (s *Socket) Close() error {
	mu.Lock()
	defer mu.Unlock()
	return s.closeLocked()
}

func (s *Socket) closeLocked() error {
	if s.closed {
		return nil
	}
	// Calling this deletes the pointer. It must not be referred to after
	// this.
	C.utp_destroy(s.ctx)
	s.ctx = nil
	s.pc.Close()
	close(s.backlog)
	close(s.nonUtpReads)
	s.closed = true
	return nil
}

func (s *Socket) Addr() net.Addr {
	return s.pc.LocalAddr()
}

func (s *Socket) LocalAddr() net.Addr {
	return s.pc.LocalAddr()
}

func (s *Socket) Accept() (net.Conn, error) {
	nc, ok := <-s.backlog
	if !ok {
		return nil, errors.New("closed")
	}
	return nc, nil
}

func (s *Socket) Dial(addr string) (net.Conn, error) {
	return s.DialTimeout(addr, 0)
}

func (s *Socket) DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
	ctx := context.Background()
	if timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel()
	}
	return s.DialContext(ctx, "", addr)
}

func (s *Socket) resolveAddr(network, addr string) (net.Addr, error) {
	if network == "" {
		network = s.Addr().Network()
	}
	return resolveAddr(network, addr)
}

func resolveAddr(network, addr string) (net.Addr, error) {
	switch network {
	case "inproc":
		return inproc.ResolveAddr(network, addr)
	default:
		return net.ResolveUDPAddr(network, addr)
	}
}

// Passing an empty network will use the network of the Socket's listener.
func (s *Socket) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
	c, err := s.NewConn()
	if err != nil {
		return nil, err
	}
	err = c.Connect(ctx, network, addr)
	if err != nil {
		c.Close()
		return nil, err
	}
	return c, nil
}

func (s *Socket) NewConn() (*Conn, error) {
	mu.Lock()
	defer mu.Unlock()
	if s.closed {
		return nil, errors.New("socket closed")
	}
	return s.newConn(C.utp_create_socket(s.ctx)), nil
}

func (s *Socket) pushBacklog(c *Conn) {
	select {
	case s.backlog <- c:
	default:
		c.close()
	}
}

func (s *Socket) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	p, ok := <-s.nonUtpReads
	if !ok {
		err = errors.New("closed")
		return
	}
	n = copy(b, p.b)
	addr = p.from
	return
}

func (s *Socket) onReadNonUtp(b []byte, from net.Addr) {
	if s.closed {
		return
	}
	socketNonUtpPacketsReceived.Add(1)
	select {
	case s.nonUtpReads <- packet{append([]byte(nil), b...), from}:
	default:
		// log.Printf("dropped non utp packet: no room in buffer")
		nonUtpPacketsDropped.Add(1)
	}
}

func (s *Socket) SetReadDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) SetWriteDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) SetDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) WriteTo(b []byte, addr net.Addr) (int, error) {
	return s.pc.WriteTo(b, addr)
}

func (s *Socket) ReadBufferLen() int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_get_option(s.ctx, C.UTP_RCVBUF))
}

func (s *Socket) WriteBufferLen() int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_get_option(s.ctx, C.UTP_SNDBUF))
}

func (s *Socket) SetWriteBufferLen(len int) {
	mu.Lock()
	defer mu.Unlock()
	i := C.utp_context_set_option(s.ctx, C.UTP_SNDBUF, C.int(len))
	if i != 0 {
		panic(i)
	}
}

func (s *Socket) SetOption(opt Option, val int) int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_set_option(s.ctx, opt, C.int(val)))
}

func (s *Socket) SetFirewallCallback(f FirewallCallback) {
	mu.Lock()
	s.firewallCallback = f
	mu.Unlock()
}
<|endoftext|>


<commit_before>package torrent

import (
	"context"
	"fmt"
	"net"
	"net/url"
	"strconv"
	"strings"

	"github.com/anacrolix/missinggo"
	"github.com/anacrolix/missinggo/perf"
	"golang.org/x/net/proxy"
)

type dialer interface {
	dial(_ context.Context, addr string) (net.Conn, error)
}

type socket interface {
	net.Listener
	dialer
}

func getProxyDialer(proxyURL string) (proxy.Dialer, error) {
	fixedURL, err := url.Parse(proxyURL)
	if err != nil {
		return nil, err
	}

	return proxy.FromURL(fixedURL, proxy.Direct)
}

func listen(network, addr, proxyURL string, f firewallCallback) (socket, error) {
	if isTcpNetwork(network) {
		return listenTcp(network, addr, proxyURL)
	} else if isUtpNetwork(network) {
		return listenUtp(network, addr, proxyURL, f)
	} else {
		panic(fmt.Sprintf("unknown network %q", network))
	}
}

func isTcpNetwork(s string) bool {
	return strings.Contains(s, "tcp")
}

func isUtpNetwork(s string) bool {
	return strings.Contains(s, "utp") || strings.Contains(s, "udp")
}

func listenTcp(network, address, proxyURL string) (s socket, err error) {
	l, err := net.Listen(network, address)
	if err != nil {
		return
	}
	defer func() {
		if err != nil {
			l.Close()
		}
	}()

	// If we don't need the proxy - then we should return default net.Dialer,
	// otherwise, let's try to parse the proxyURL and return proxy.Dialer
	if len(proxyURL) != 0 {
		// TODO: The error should be propagated, as proxy may be in use for
		// security or privacy reasons. Also just pass proxy.Dialer in from
		// the Config.
		if dialer, err := getProxyDialer(proxyURL); err == nil {
			return tcpSocket{l, func(ctx context.Context, addr string) (conn net.Conn, err error) {
				defer perf.ScopeTimerErr(&err)()
				return dialer.Dial(network, addr)
			}}, nil
		}
	}
	dialer := net.Dialer{}
	return tcpSocket{l, func(ctx context.Context, addr string) (conn net.Conn, err error) {
		defer perf.ScopeTimerErr(&err)()
		return dialer.DialContext(ctx, network, addr)
	}}, nil
}

type tcpSocket struct {
	net.Listener
	d func(ctx context.Context, addr string) (net.Conn, error)
}

func (me tcpSocket) dial(ctx context.Context, addr string) (net.Conn, error) {
	return me.d(ctx, addr)
}

func setPort(addr string, port int) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		panic(err)
	}
	return net.JoinHostPort(host, strconv.FormatInt(int64(port), 10))
}

func listenAll(networks []string, getHost func(string) string, port int, proxyURL string, f firewallCallback) ([]socket, error) {
	if len(networks) == 0 {
		return nil, nil
	}
	var nahs []networkAndHost
	for _, n := range networks {
		nahs = append(nahs, networkAndHost{n, getHost(n)})
	}
	for {
		ss, retry, err := listenAllRetry(nahs, port, proxyURL, f)
		if !retry {
			return ss, err
		}
	}
}

type networkAndHost struct {
	Network string
	Host    string
}

func listenAllRetry(nahs []networkAndHost, port int, proxyURL string, f firewallCallback) (ss []socket, retry bool, err error) {
	ss = make([]socket, 1, len(nahs))
	portStr := strconv.FormatInt(int64(port), 10)
	ss[0], err = listen(nahs[0].Network, net.JoinHostPort(nahs[0].Host, portStr), proxyURL, f)
	if err != nil {
		return nil, false, fmt.Errorf("first listen: %s", err)
	}
	defer func() {
		if err != nil || retry {
			for _, s := range ss {
				s.Close()
			}
			ss = nil
		}
	}()
	portStr = strconv.FormatInt(int64(missinggo.AddrPort(ss[0].Addr())), 10)
	for _, nah := range nahs[1:] {
		s, err := listen(nah.Network, net.JoinHostPort(nah.Host, portStr), proxyURL, f)
		if err != nil {
			return ss,
				missinggo.IsAddrInUse(err) && port == 0,
				fmt.Errorf("subsequent listen: %s", err)
		}
		ss = append(ss, s)
	}
	return
}

type firewallCallback func(net.Addr) bool

func listenUtp(network, addr, proxyURL string, fc firewallCallback) (s socket, err error) {
	us, err := NewUtpSocket(network, addr, fc)
	if err != nil {
		return
	}

	// If we don't need the proxy - then we should return default net.Dialer,
	// otherwise, let's try to parse the proxyURL and return proxy.Dialer
	if len(proxyURL) != 0 {
		if dialer, err := getProxyDialer(proxyURL); err == nil {
			return utpSocketSocket{us, network, dialer}, nil
		}
	}

	return utpSocketSocket{us, network, nil}, nil
}

type utpSocketSocket struct {
	utpSocket
	network string
	d       proxy.Dialer
}

func (me utpSocketSocket) dial(ctx context.Context, addr string) (conn net.Conn, err error) {
	defer perf.ScopeTimerErr(&err)()
	if me.d != nil {
		return me.d.Dial(me.network, addr)
	}

	return me.utpSocket.DialContext(ctx, me.network, addr)
}
<commit_msg>Improve error handling for listening<commit_after>package torrent

import (
	"context"
	"fmt"
	"net"
	"net/url"
	"strconv"
	"strings"

	"github.com/anacrolix/missinggo"
	"github.com/anacrolix/missinggo/perf"
	"github.com/pkg/errors"
	"golang.org/x/net/proxy"
)

type dialer interface {
	dial(_ context.Context, addr string) (net.Conn, error)
}

type socket interface {
	net.Listener
	dialer
}

func getProxyDialer(proxyURL string) (proxy.Dialer, error) {
	fixedURL, err := url.Parse(proxyURL)
	if err != nil {
		return nil, err
	}

	return proxy.FromURL(fixedURL, proxy.Direct)
}

func listen(network, addr, proxyURL string, f firewallCallback) (socket, error) {
	if isTcpNetwork(network) {
		return listenTcp(network, addr, proxyURL)
	} else if isUtpNetwork(network) {
		return listenUtp(network, addr, proxyURL, f)
	} else {
		panic(fmt.Sprintf("unknown network %q", network))
	}
}

func isTcpNetwork(s string) bool {
	return strings.Contains(s, "tcp")
}

func isUtpNetwork(s string) bool {
	return strings.Contains(s, "utp") || strings.Contains(s, "udp")
}

func listenTcp(network, address, proxyURL string) (s socket, err error) {
	l, err := net.Listen(network, address)
	if err != nil {
		return
	}
	defer func() {
		if err != nil {
			l.Close()
		}
	}()

	// If we don't need the proxy - then we should return default net.Dialer,
	// otherwise, let's try to parse the proxyURL and return proxy.Dialer
	if len(proxyURL) != 0 {
		// TODO: The error should be propagated, as proxy may be in use for
		// security or privacy reasons. Also just pass proxy.Dialer in from
		// the Config.
		if dialer, err := getProxyDialer(proxyURL); err == nil {
			return tcpSocket{l, func(ctx context.Context, addr string) (conn net.Conn, err error) {
				defer perf.ScopeTimerErr(&err)()
				return dialer.Dial(network, addr)
			}}, nil
		}
	}
	dialer := net.Dialer{}
	return tcpSocket{l, func(ctx context.Context, addr string) (conn net.Conn, err error) {
		defer perf.ScopeTimerErr(&err)()
		return dialer.DialContext(ctx, network, addr)
	}}, nil
}

type tcpSocket struct {
	net.Listener
	d func(ctx context.Context, addr string) (net.Conn, error)
}

func (me tcpSocket) dial(ctx context.Context, addr string) (net.Conn, error) {
	return me.d(ctx, addr)
}

func setPort(addr string, port int) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		panic(err)
	}
	return net.JoinHostPort(host, strconv.FormatInt(int64(port), 10))
}

func listenAll(networks []string, getHost func(string) string, port int, proxyURL string, f firewallCallback) ([]socket, error) {
	if len(networks) == 0 {
		return nil, nil
	}
	var nahs []networkAndHost
	for _, n := range networks {
		nahs = append(nahs, networkAndHost{n, getHost(n)})
	}
	for {
		ss, retry, err := listenAllRetry(nahs, port, proxyURL, f)
		if !retry {
			return ss, err
		}
	}
}

type networkAndHost struct {
	Network string
	Host    string
}

func listenAllRetry(nahs []networkAndHost, port int, proxyURL string, f firewallCallback) (ss []socket, retry bool, err error) {
	ss = make([]socket, 1, len(nahs))
	portStr := strconv.FormatInt(int64(port), 10)
	ss[0], err = listen(nahs[0].Network, net.JoinHostPort(nahs[0].Host, portStr), proxyURL, f)
	if err != nil {
		return nil, false, errors.Wrap(err, "first listen")
	}
	defer func() {
		if err != nil || retry {
			for _, s := range ss {
				s.Close()
			}
			ss = nil
		}
	}()
	portStr = strconv.FormatInt(int64(missinggo.AddrPort(ss[0].Addr())), 10)
	for _, nah := range nahs[1:] {
		s, err := listen(nah.Network, net.JoinHostPort(nah.Host, portStr), proxyURL, f)
		if err != nil {
			return ss,
				missinggo.IsAddrInUse(err) && port == 0,
				errors.Wrap(err, "subsequent listen")
		}
		ss = append(ss, s)
	}
	return
}

type firewallCallback func(net.Addr) bool

func listenUtp(network, addr, proxyURL string, fc firewallCallback) (s socket, err error) {
	us, err := NewUtpSocket(network, addr, fc)
	if err != nil {
		return
	}

	// If we don't need the proxy - then we should return default net.Dialer,
	// otherwise, let's try to parse the proxyURL and return proxy.Dialer
	if len(proxyURL) != 0 {
		if dialer, err := getProxyDialer(proxyURL); err == nil {
			return utpSocketSocket{us, network, dialer}, nil
		}
	}

	return utpSocketSocket{us, network, nil}, nil
}

type utpSocketSocket struct {
	utpSocket
	network string
	d       proxy.Dialer
}

func (me utpSocketSocket) dial(ctx context.Context, addr string) (conn net.Conn, err error) {
	defer perf.ScopeTimerErr(&err)()
	if me.d != nil {
		return me.d.Dial(me.network, addr)
	}

	return me.utpSocket.DialContext(ctx, me.network, addr)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"sync"

	"github.com/syndtr/goleveldb/leveldb"
)

var (
	// 因为go的正则在匹配ReadAll出来的[]byte时有点奇怪，所以必须用[^\w]{0,}?来代替\n

	// 用于获取版块名称
	getForumName = regexp.MustCompile(`<h1 class="xs2">[^\w]{0,}?<a[^>]*>([^<]*)</a>`)
	// 用于获取版块帖子分页数量
	getForumPageNumber = regexp.MustCompile(`<a href="[^"]+" class="last">\.\.\. ([0-9]+)</a>`)
	// 用于获取版块介绍
	getForumIntroduction = regexp.MustCompile(`<div id="forum_rules_[0-9]*"[^>]*>([\w\W]{0,}?)(?:</div>[^\w]{0,}?){3}<div class="drag">`)
	// 用于获取帖子信息
	getPostInfo = regexp.MustCompile(`
\n
\n\n(?:<[^>]+>\n)+	\n\n\[]*>([^<]+)\]\s*]*>([^<]+)\n(?:<[^\n]+>\n){0,}?	\n\n\n]*>([^<]+)\n(?:)?([^<]+)(){1,2}\n	\n(?:<[^\n]+>\n){7}`)
	// 用于获取帖子内容
	getPostBody = regexp.MustCompile(`<div class="pcb">((?:.*\n){0,}?)<div id="comment_[0-9]*" class="cm">`)
)

func main() {
	if len(os.Args) == 1 {
		fmt.Println("请在命令后添加论坛版块FID参数")
		fmt.Println("FID可从版块URL中寻找")
		fmt.Println("数据会保存在db/FID路径中")
		return
	}

	// 论坛版块fid
	fid := os.Args[1]

	resp, err := http.Get("http://www.mcbbs.net/forum.php?mod=forumdisplay&fid=" + fid + "&orderby=dateline&page=1")
	if err != nil {
		printError("http.Get", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		printError("ReadAll", err)
		os.Exit(1)
	}

	if bytes.Contains(body, []byte(`<p>抱歉，指定的版块不存在</p>`)) {
		printError("GetForum", "版块不存在，请检查PID是否正确")
		os.Exit(1)
	}

	// 获取版块名称
	n := getForumName.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumName", "获取版块名称出错")
		os.Exit(1)
	}
	forumName := string(n[1])
	printInfo("版块名称", forumName)

	// 获取版块帖子分页数量
	n = getForumPageNumber.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumPageNumber", "获取版块分页数量出错")
		os.Exit(1)
	}
	maxPagesNum, err := strconv.Atoi(string(n[1]))
	if err != nil {
		printError("Atoi", err)
		os.Exit(1)
	}
	printInfo("本版块全部分页数量", maxPagesNum)

	// 获取版块介绍
	n = getForumIntroduction.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumIntroduction", "获取版块介绍出错")
		os.Exit(1)
	}
	forumIntroduction := string(n[1])
	print(forumIntroduction)

	// 创建数据库
	db, err := leveldb.OpenFile("db/"+fid, nil)
	if err != nil {
		printError("OpenDB", err)
		os.Exit(1)
	}
	defer db.Close()

	// 保存版块信息到数据库
	buf, err := encode(&forumInfo{
		Name:         forumName,
		PageNumber:   maxPagesNum,
		Introduction: forumIntroduction,
	})
	if err != nil {
		printError("EncodeForumInfo", err)
	}
	err = db.Put([]byte("info"), buf, nil)
	if err != nil {
		printError("PutForumInfo", err)
		os.Exit(1)
	}

	// 用于等待全部线程执行完毕
	var wg sync.WaitGroup
	// 用于统计还有多少页未完成
	var pageAmount int
	// 获取每一页的所有帖子
	for i := 0; i < maxPagesNum; i++ {
		wg.Add(1)
		pageAmount++
		go func(page int) {
			postList, err := getPagesList(fid, page)
			if err != nil {
				printError("getPagesList"+string(page), err)
			} else {
				// 获取每个帖子的内容
				for i, v := range postList {
					// for用于重试
					for {
						resp, err := http.Get("http://www.mcbbs.net/" + v.Url)
						if err != nil {
							printError("GetPost", err)
						} else {
							defer resp.Body.Close()

							// 检查服务器是否返回成功
							if resp.StatusCode != 200 {
								// 服务器错误
								printError("GetPost.ServerError", "服务器错误，错误码：", resp.StatusCode)
								if resp.StatusCode == 404 {
									// 如果为404，则没有重试的必要，跳出重试
									printError("GetPost.Retry", "帖子《"+v.Title+"》不存在")
									break
								}
							} else {
								body, err := ioutil.ReadAll(resp.Body)
								if err != nil {
									printError("GetPost.ReadAll", err)
								} else {
									n := getPostBody.FindSubmatch(body)
									// 检查是否获取body成功
									if len(n) == 0 {
										printError("GetPost.FindSubmatch", "未找到页面内文章部分")
									} else {
										postBody := string(n[1])
										// 存入Body
										postList[i].Body = postBody
										printInfo("GetPost", v.Title, "[OK]")
										// 跳出循环重试
										break
									}
								}
							}
						}
						// 获取失败，重试
						printError("GetPost.Retry", "获取帖子《"+v.Title+"》失败，正在重试")
					}
				}
			}
			// 编码
			byt, err := encode(postList)
			if err != nil {
				printError("Encode", err)
				return
			}
			// 存入数据库
			err = db.Put([]byte("page_"+strconv.Itoa(page)), byt, nil)
			if err != nil {
				printError("db.Put", err)
				return
			}

			pageAmount--

			printInfo("OK", "线程", page, "执行完毕")
			printInfo("OK", "版块分页", page, "中的所有帖子已经储存到本地，还有", pageAmount, "页正在下载中")
			wg.Done()
		}(i + 1)
	}
	wg.Wait()
	printInfo("OK", "FID为", fid, "的版块中的所有帖子已储存到本地")
}

// 获取单页面的帖子列表
func getPagesList(fid string, pageNum int) (pageList []postInfo, err error) {
	for {
		resp, err := http.Get("http://www.mcbbs.net/forum.php?mod=forumdisplay&fid=" + fid + "&orderby=dateline&page=" + strconv.Itoa(pageNum))
		if err != nil {
			printError("GetPageList", err)
		} else {
			defer resp.Body.Close()

			// 检查服务器是否返回成功
			if resp.StatusCode != 200 {
				// 服务器错误
				printError("GetPost.ServerError", "服务器错误，错误码：", resp.StatusCode)
			} else {
				body, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					printError("GetPageList.ReadAll", err)
				} else {
					/*
						// getPage所匹配的信息
						// 获取到的参数已用“m[x]说明”标出

						[m[1]分类] m[3]标题

						 ...234
						New

						m[4]作者
						m[6]发帖时间2/span>

						111189

						谢普
						5 天前

					*/
					m := getPostInfo.FindAllStringSubmatch(string(body), -1)

					// 检查是否获取成功
					if len(m) == 0 {
						printError("GetPageList.FindAllStringSubmatch", "未找到分页内帖子")
					} else {
						// 处理单个帖子信息
						for _, v := range m {
							date := v[5]
							// 如果发帖时间1为空，那么使用发帖时间2
							if date == "" {
								date = v[6]
							}
							// 储存帖子信息
							postInf := postInfo{
								Category: v[1],
								Url:      v[2],
								Title:    v[3],
								Author:   v[4],
								Date:     date,
							}
							pageList = append(pageList, postInf)
						}
						// 跳出重试
						break
					}
				}
			}
		}
		// 获取失败，重试
		printError("getPageList", "获取版块分页", pageNum, "失败，正在重试")
	}

	return pageList, nil
}

func printInfo(s string, v ...interface{}) {
	log.Println(append([]interface{}{"[INFO]", s + ":"}, v...)...)
}

func printError(s string, v ...interface{}) {
	log.Println(append([]interface{}{"[ERROR]", s + ":"}, v...)...)
}

func encode(data interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(data)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decode(data []byte, to interface{}) error {
	buf := bytes.NewBuffer(data)
	dec := gob.NewDecoder(buf)
	return dec.Decode(to)
}

type forumInfo struct {
	Name         string
	PageNumber   int
	Introduction string
}

type postInfo struct {
	Category string
	Url      string
	Title    string
	Author   string
	Date     string
	Body     string
}
<commit_msg>modify: 部分正则细节<commit_after>package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"sync"

	"github.com/syndtr/goleveldb/leveldb"
)

var (
	// go的正则在匹配ReadAll出来的数据时有点奇怪，必须用[\w\W]{0,}?来代替\n

	// 用于获取版块名称
	getForumName = regexp.MustCompile(`<h1 class="xs2">[\w\W]{0,}?<a[^>]*>([^<]*)</a>`)
	// 用于获取版块帖子分页数量
	getForumPageNumber = regexp.MustCompile(`<a href="[^"]+" class="last">\.\.\. ([0-9]+)</a>`)
	// 用于获取版块介绍
	getForumIntroduction = regexp.MustCompile(`<div id="forum_rules_[0-9]*"[^>]*>([\w\W]{0,}?)(?:</div>[\w\W]{0,}?){3}<div class="drag">`)
	// 用于获取帖子信息
	getPostInfo = regexp.MustCompile(`
\n
\n\n(?:<[^>]+>\n)+	\n\n\[]*>([^<]+)\]\s*]*>([^<]+)\n(?:<[^\n]+>\n){0,}?	\n\n\n]*>([^<]+)\n(?:)?([^<]+)(){1,2}\n	\n(?:<[^\n]+>\n){7}`)
	// 用于获取帖子内容
	getPostBody = regexp.MustCompile(`<div class="pcb">((?:.*\n){0,}?)<div id="comment_[0-9]*" class="cm">`)
)

func main() {
	if len(os.Args) == 1 {
		fmt.Println("请在命令后添加论坛版块FID参数")
		fmt.Println("FID可从版块URL中寻找")
		fmt.Println("数据会保存在db/FID路径中")
		return
	}

	// 论坛版块fid
	fid := os.Args[1]

	resp, err := http.Get("http://www.mcbbs.net/forum.php?mod=forumdisplay&fid=" + fid + "&orderby=dateline&page=1")
	if err != nil {
		printError("http.Get", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		printError("ReadAll", err)
		os.Exit(1)
	}

	if bytes.Contains(body, []byte(`<p>抱歉，指定的版块不存在</p>`)) {
		printError("GetForum", "版块不存在，请检查PID是否正确")
		os.Exit(1)
	}

	// 获取版块名称
	n := getForumName.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumName", "获取版块名称出错")
		os.Exit(1)
	}
	forumName := string(n[1])
	printInfo("版块名称", forumName)

	// 获取版块帖子分页数量
	n = getForumPageNumber.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumPageNumber", "获取版块分页数量出错")
		os.Exit(1)
	}
	maxPagesNum, err := strconv.Atoi(string(n[1]))
	if err != nil {
		printError("Atoi", err)
		os.Exit(1)
	}
	printInfo("本版块全部分页数量", maxPagesNum)

	// 获取版块介绍
	n = getForumIntroduction.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumIntroduction", "获取版块介绍出错")
		os.Exit(1)
	}
	forumIntroduction := string(n[1])
	print(forumIntroduction)

	// 创建数据库
	db, err := leveldb.OpenFile("db/"+fid, nil)
	if err != nil {
		printError("OpenDB", err)
		os.Exit(1)
	}
	defer db.Close()

	// 保存版块信息到数据库
	buf, err := encode(&forumInfo{
		Name:         forumName,
		PageNumber:   maxPagesNum,
		Introduction: forumIntroduction,
	})
	if err != nil {
		printError("EncodeForumInfo", err)
	}
	err = db.Put([]byte("info"), buf, nil)
	if err != nil {
		printError("PutForumInfo", err)
		os.Exit(1)
	}

	// 用于等待全部线程执行完毕
	var wg sync.WaitGroup
	// 用于统计还有多少页未完成
	var pageAmount int
	// 获取每一页的所有帖子
	for i := 0; i < maxPagesNum; i++ {
		wg.Add(1)
		pageAmount++
		go func(page int) {
			postList, err := getPagesList(fid, page)
			if err != nil {
				printError("getPagesList"+string(page), err)
			} else {
				// 获取每个帖子的内容
				for i, v := range postList {
					// for用于重试
					for {
						resp, err := http.Get("http://www.mcbbs.net/" + v.Url)
						if err != nil {
							printError("GetPost", err)
						} else {
							defer resp.Body.Close()

							// 检查服务器是否返回成功
							if resp.StatusCode != 200 {
								// 服务器错误
								printError("GetPost.ServerError", "服务器错误，错误码：", resp.StatusCode)
								if resp.StatusCode == 404 {
									// 如果为404，则没有重试的必要，跳出重试
									printError("GetPost.Retry", "帖子《"+v.Title+"》不存在")
									break
								}
							} else {
								body, err := ioutil.ReadAll(resp.Body)
								if err != nil {
									printError("GetPost.ReadAll", err)
								} else {
									n := getPostBody.FindSubmatch(body)
									// 检查是否获取body成功
									if len(n) == 0 {
										printError("GetPost.FindSubmatch", "未找到页面内文章部分")
									} else {
										postBody := string(n[1])
										// 存入Body
										postList[i].Body = postBody
										printInfo("GetPost", v.Title, "[OK]")
										// 跳出循环重试
										break
									}
								}
							}
						}
						// 获取失败，重试
						printError("GetPost.Retry", "获取帖子《"+v.Title+"》失败，正在重试")
					}
				}
			}
			// 编码
			byt, err := encode(postList)
			if err != nil {
				printError("Encode", err)
				return
			}
			// 存入数据库
			err = db.Put([]byte("page_"+strconv.Itoa(page)), byt, nil)
			if err != nil {
				printError("db.Put", err)
				return
			}

			pageAmount--

			printInfo("OK", "线程", page, "执行完毕")
			printInfo("OK", "版块分页", page, "中的所有帖子已经储存到本地，还有", pageAmount, "页正在下载中")
			wg.Done()
		}(i + 1)
	}
	wg.Wait()
	printInfo("OK", "FID为", fid, "的版块中的所有帖子已储存到本地")
}

// 获取单页面的帖子列表
func getPagesList(fid string, pageNum int) (pageList []postInfo, err error) {
	for {
		resp, err := http.Get("http://www.mcbbs.net/forum.php?mod=forumdisplay&fid=" + fid + "&orderby=dateline&page=" + strconv.Itoa(pageNum))
		if err != nil {
			printError("GetPageList", err)
		} else {
			defer resp.Body.Close()

			// 检查服务器是否返回成功
			if resp.StatusCode != 200 {
				// 服务器错误
				printError("GetPost.ServerError", "服务器错误，错误码：", resp.StatusCode)
			} else {
				body, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					printError("GetPageList.ReadAll", err)
				} else {
					/*
						// getPage所匹配的信息
						// 获取到的参数已用“m[x]说明”标出

						[m[1]分类] m[3]标题

						 ...234
						New

						m[4]作者
						m[6]发帖时间2/span>

						111189

						谢普
						5 天前

					*/
					m := getPostInfo.FindAllStringSubmatch(string(body), -1)

					// 检查是否获取成功
					if len(m) == 0 {
						printError("GetPageList.FindAllStringSubmatch", "未找到分页内帖子")
					} else {
						// 处理单个帖子信息
						for _, v := range m {
							date := v[5]
							// 如果发帖时间1为空，那么使用发帖时间2
							if date == "" {
								date = v[6]
							}
							// 储存帖子信息
							postInf := postInfo{
								Category: v[1],
								Url:      v[2],
								Title:    v[3],
								Author:   v[4],
								Date:     date,
							}
							pageList = append(pageList, postInf)
						}
						// 跳出重试
						break
					}
				}
			}
		}
		// 获取失败，重试
		printError("getPageList", "获取版块分页", pageNum, "失败，正在重试")
	}

	return pageList, nil
}

func printInfo(s string, v ...interface{}) {
	log.Println(append([]interface{}{"[INFO]", s + ":"}, v...)...)
}

func printError(s string, v ...interface{}) {
	log.Println(append([]interface{}{"[ERROR]", s + ":"}, v...)...)
}

func encode(data interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(data)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decode(data []byte, to interface{}) error {
	buf := bytes.NewBuffer(data)
	dec := gob.NewDecoder(buf)
	return dec.Decode(to)
}

type forumInfo struct {
	Name         string
	PageNumber   int
	Introduction string
}

type postInfo struct {
	Category string
	Url      string
	Title    string
	Author   string
	Date     string
	Body     string
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/DeedleFake/census/ps2/events"
	"log"
	"os"
	"os/signal"
	"time"
)

var (
	// session can be used to get a copy of the current session.
	session = make(chan Session)
)

// coord coordinates the session, updating it properly when login and
// logout events occur, and sending a copy of the session down the
// session channel when it's requested.
func coord(logins <-chan *events.PlayerLogin, logouts <-chan *events.PlayerLogout, errors <-chan error) {
	log.Printf("Loading session from %q...", flags.session)
	s, err := LoadSession(flags.session)
	if err != nil {
		log.Printf("Failed to load session from %q: %v", flags.session, err)
		log.Println("Creating new session...")
	}
	s.Runtime = timeDiff(time.Now())
	s.Longest = 0
	s.Shortest = jsonDuration(1000 * time.Hour) // Just need something ridiculous.

	chars := make(map[int64]time.Time)

	copySession := func() Session {
		s := s
		s.NumChars = len(chars)

		return s
	}

	for {
		select {
		case ev := <-logins:
			chars[ev.CharacterID] = time.Unix(ev.Timestamp, 0)
		case ev := <-logouts:
			// TODO: Use the REST API to get login times, rather than
			// tracking it manually. I don't know why I didn't do that in
			// the first place...
			if in, ok := chars[ev.CharacterID]; ok {
				d := time.Unix(ev.Timestamp, 0).Sub(in)

				s.Total.Update(d)
				if d > flags.short {
					s.NoShort.Update(d)
				}

				if d > time.Duration(s.Longest) {
					s.Longest = jsonDuration(d)
				}
				if d < time.Duration(s.Shortest) {
					s.Shortest = jsonDuration(d)
				}

				delete(chars, ev.CharacterID)
			}

		case err := <-errors:
			s.Err = err

		case session <- copySession():
		}
	}
}

// monitor connects to the census API, subscribes to PlayerLogin and
// PlayerLogout events, and then sends them down the appropriate
// channels.
func monitor(logins chan<- *events.PlayerLogin, logouts chan<- *events.PlayerLogout, errors chan<- error) {
	cl, err := events.NewClient("", "", "example")
	if err != nil {
		log.Fatalf("Failed to open client: %v", err)
	}
	defer cl.Close()

	err = cl.Subscribe(events.Sub{
		Events: []string{"PlayerLogin", "PlayerLogout"},
		Chars:  events.SubAll,
		Worlds: events.SubAll,
	})
	if err != nil {
		log.Fatalf("Failed to subscribe to login/logout events: %v", err)
	}

	for {
		ev, err := cl.Next()
		if err != nil {
			log.Printf("Error while fetching event: %v", err)
			errors <- err
			continue
		}
		errors <- nil

		switch ev := ev.(type) {
		case *events.PlayerLogin:
			logins <- ev
		case *events.PlayerLogout:
			logouts <- ev
		}
	}
}

func main() {
	logins := make(chan *events.PlayerLogin)
	logouts := make(chan *events.PlayerLogout)
	errors := make(chan error)

	go monitor(logins, logouts, errors)
	go coord(logins, logouts, errors)
	go server()

	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt)
	<-sig

	log.Printf("Saving session to %q...", flags.session)
	err := (<-session).Save(flags.session)
	if err != nil {
		log.Printf("Failed to save session to %q: %v", flags.session, err)
	}
}
<commit_msg>Exit cleanly on SIGTERM.<commit_after>package main

import (
	"github.com/DeedleFake/census/ps2/events"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	// session can be used to get a copy of the current session.
	session = make(chan Session)
)

// coord coordinates the session, updating it properly when login and
// logout events occur, and sending a copy of the session down the
// session channel when it's requested.
func coord(logins <-chan *events.PlayerLogin, logouts <-chan *events.PlayerLogout, errors <-chan error) {
	log.Printf("Loading session from %q...", flags.session)
	s, err := LoadSession(flags.session)
	if err != nil {
		log.Printf("Failed to load session from %q: %v", flags.session, err)
		log.Println("Creating new session...")
	}
	s.Runtime = timeDiff(time.Now())
	s.Longest = 0
	s.Shortest = jsonDuration(1000 * time.Hour) // Just need something ridiculous.

	chars := make(map[int64]time.Time)

	copySession := func() Session {
		s := s
		s.NumChars = len(chars)

		return s
	}

	for {
		select {
		case ev := <-logins:
			chars[ev.CharacterID] = time.Unix(ev.Timestamp, 0)
		case ev := <-logouts:
			// TODO: Use the REST API to get login times, rather than
			// tracking it manually. I don't know why I didn't do that in
			// the first place...
			if in, ok := chars[ev.CharacterID]; ok {
				d := time.Unix(ev.Timestamp, 0).Sub(in)

				s.Total.Update(d)
				if d > flags.short {
					s.NoShort.Update(d)
				}

				if d > time.Duration(s.Longest) {
					s.Longest = jsonDuration(d)
				}
				if d < time.Duration(s.Shortest) {
					s.Shortest = jsonDuration(d)
				}

				delete(chars, ev.CharacterID)
			}

		case err := <-errors:
			s.Err = err

		case session <- copySession():
		}
	}
}

// monitor connects to the census API, subscribes to PlayerLogin and
// PlayerLogout events, and then sends them down the appropriate
// channels.
func monitor(logins chan<- *events.PlayerLogin, logouts chan<- *events.PlayerLogout, errors chan<- error) {
	cl, err := events.NewClient("", "", "example")
	if err != nil {
		log.Fatalf("Failed to open client: %v", err)
	}
	defer cl.Close()

	err = cl.Subscribe(events.Sub{
		Events: []string{"PlayerLogin", "PlayerLogout"},
		Chars:  events.SubAll,
		Worlds: events.SubAll,
	})
	if err != nil {
		log.Fatalf("Failed to subscribe to login/logout events: %v", err)
	}

	for {
		ev, err := cl.Next()
		if err != nil {
			log.Printf("Error while fetching event: %v", err)
			errors <- err
			continue
		}
		errors <- nil

		switch ev := ev.(type) {
		case *events.PlayerLogin:
			logins <- ev
		case *events.PlayerLogout:
			logouts <- ev
		}
	}
}

func main() {
	logins := make(chan *events.PlayerLogin)
	logouts := make(chan *events.PlayerLogout)
	errors := make(chan error)

	go monitor(logins, logouts, errors)
	go coord(logins, logouts, errors)
	go server()

	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	log.Printf("Caught signal %q", <-sig)

	log.Printf("Saving session to %q...", flags.session)
	err := (<-session).Save(flags.session)
	if err != nil {
		log.Printf("Failed to save session to %q: %v", flags.session, err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/BurntSushi/xgb"
	"github.com/BurntSushi/xgb/screensaver"
	"github.com/BurntSushi/xgb/xproto"
	"log"
	"time"
)

type tracker map[window]*track

type track struct {
	Start time.Time
	Spent time.Duration
}

type window struct {
	Class string
	Name  string
}

func (t track) String() string {
	return fmt.Sprint(t.Spent)
}

func (w window) String() string {
	return fmt.Sprintf("%s: %s", w.Class, w.Name)
}

func getClass(b []byte) string {
	i := bytes.IndexByte(b, 0)
	if i == -1 {
		return ""
	}
	return string(b[:i])
}

func asciizToString(b []byte) (s []string) {
	for _, x := range bytes.Split(b, []byte{0}) {
		s = append(s, string(x))
	}
	if len(s) > 0 && s[len(s)-1] == "" {
		s = s[:len(s)-1]
	}
	return
}

func atom(X *xgb.Conn, aname string) *xproto.InternAtomReply {
	a, err := xproto.InternAtom(X, true, uint16(len(aname)), aname).Reply()
	if err != nil {
		log.Fatal("atom", err)
	}
	return a
}

func prop(X *xgb.Conn, w xproto.Window, a *xproto.InternAtomReply) *xproto.GetPropertyReply {
	p, err := xproto.GetProperty(X, false, w, a.Atom, xproto.GetPropertyTypeAny, 0, (1<<32)-1).Reply()
	if err != nil {
		log.Fatal("property", err)
	}
	return p
}

func winName(X *xgb.Conn, root xproto.Window) (window, bool) {
	activeAtom := atom(X, "_NET_ACTIVE_WINDOW")
	netNameAtom := atom(X, "_NET_WM_NAME")
	nameAtom := atom(X, "WM_NAME")
	classAtom := atom(X, "WM_CLASS")

	active := prop(X, root, activeAtom)
	windowId := xproto.Window(xgb.Get32(active.Value))

	/* skip root window */
	if windowId == 0 {
		return window{}, false
	}

	spy(X, windowId)

	name := prop(X, windowId, netNameAtom)
	if string(name.Value) == "" {
		name = prop(X, windowId, nameAtom)
	}
	class := prop(X, windowId, classAtom)

	w := window{
		Class: getClass(class.Value),
		Name:  string(name.Value),
	}

	return w, true
}

func rootWin(X *xgb.Conn) xproto.Window {
	setup := xproto.Setup(X)
	return setup.DefaultScreen(X).Root
}

func spy(X *xgb.Conn, w xproto.Window) {
	xproto.ChangeWindowAttributes(X, w, xproto.CwEventMask,
		[]uint32{xproto.EventMaskPropertyChange})
}

func collect(tracks tracker) {
	var prev *track

	X, err := xgb.NewConn()
	if err != nil {
		log.Fatal("xgb", err)
	}
	defer X.Close()

	err = screensaver.Init(X)
	if err != nil {
		log.Fatal("screensaver", err)
	}

	root := rootWin(X)
	spy(X, root)
	drw := xproto.Drawable(root)
	screensaver.SelectInput(X, drw, screensaver.EventNotifyMask)

	for {
		ev, everr := X.WaitForEvent()
		if everr != nil {
			log.Fatal("wait for event", everr)
		}
		switch event := ev.(type) {
		case xproto.PropertyNotifyEvent:
			if prev != nil {
				prev.Spent += time.Since(prev.Start)
			}
			if win, ok := winName(X, root); ok {
				if _, ok := tracks[win]; !ok {
					tracks[win] = new(track)
				}
				tracks[win].Start = time.Now()
				prev = tracks[win]
			}
		case screensaver.NotifyEvent:
			switch event.State {
			case screensaver.StateOn:
				fmt.Println("screensaver on")
				prev = nil
			}
		}
	}
}

func display(tracks tracker) {
	for {
		var total time.Duration
		classtotal := make(map[string]time.Duration)
		for n, t := range tracks {
			fmt.Println(n, t)
			total += t.Spent
			classtotal[n.Class] += t.Spent
		}
		fmt.Println("")
		for k, v := range classtotal {
			fmt.Println(k, v)
		}
		fmt.Println("Total:", total)
		fmt.Println("")
		time.Sleep(3 * time.Second)
	}
}

func cleanup(tracks tracker) {
	for {
		for k, v := range tracks {
			if time.Since(v.Start).Hours() > 12.0 {
				log.Println("removing", k)
				delete(tracks, k)
			}
		}
		time.Sleep(time.Minute)
	}
}

func main() {
	tracks := make(tracker)
	go collect(tracks)
	go cleanup(tracks)
	display(tracks)
}
<commit_msg>count first window too<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/BurntSushi/xgb"
	"github.com/BurntSushi/xgb/screensaver"
	"github.com/BurntSushi/xgb/xproto"
	"log"
	"time"
)

type tracker map[window]*track

type track struct {
	Start time.Time
	Spent time.Duration
}

type window struct {
	Class string
	Name  string
}

func (t track) String() string {
	return fmt.Sprint(t.Spent)
}

func (w window) String() string {
	return fmt.Sprintf("%s: %s", w.Class, w.Name)
}

func getClass(b []byte) string {
	i := bytes.IndexByte(b, 0)
	if i == -1 {
		return ""
	}
	return string(b[:i])
}

func asciizToString(b []byte) (s []string) {
	for _, x := range bytes.Split(b, []byte{0}) {
		s = append(s, string(x))
	}
	if len(s) > 0 && s[len(s)-1] == "" {
		s = s[:len(s)-1]
	}
	return
}

func atom(X *xgb.Conn, aname string) *xproto.InternAtomReply {
	a, err := xproto.InternAtom(X, true, uint16(len(aname)), aname).Reply()
	if err != nil {
		log.Fatal("atom", err)
	}
	return a
}

func prop(X *xgb.Conn, w xproto.Window, a *xproto.InternAtomReply) *xproto.GetPropertyReply {
	p, err := xproto.GetProperty(X, false, w, a.Atom, xproto.GetPropertyTypeAny, 0, (1<<32)-1).Reply()
	if err != nil {
		log.Fatal("property", err)
	}
	return p
}

func winName(X *xgb.Conn, root xproto.Window) (window, bool) {
	activeAtom := atom(X, "_NET_ACTIVE_WINDOW")
	netNameAtom := atom(X, "_NET_WM_NAME")
	nameAtom := atom(X, "WM_NAME")
	classAtom := atom(X, "WM_CLASS")

	active := prop(X, root, activeAtom)
	windowId := xproto.Window(xgb.Get32(active.Value))

	/* skip root window */
	if windowId == 0 {
		return window{}, false
	}

	spy(X, windowId)

	name := prop(X, windowId, netNameAtom)
	if string(name.Value) == "" {
		name = prop(X, windowId, nameAtom)
	}
	class := prop(X, windowId, classAtom)

	w := window{
		Class: getClass(class.Value),
		Name:  string(name.Value),
	}

	return w, true
}

func rootWin(X *xgb.Conn) xproto.Window {
	setup := xproto.Setup(X)
	return setup.DefaultScreen(X).Root
}

func spy(X *xgb.Conn, w xproto.Window) {
	xproto.ChangeWindowAttributes(X, w, xproto.CwEventMask,
		[]uint32{xproto.EventMaskPropertyChange})
}

func (t tracker) Update(X *xgb.Conn, w xproto.Window) (prev *track) {
	if win, ok := winName(X, w); ok {
		if _, ok := t[win]; !ok {
			t[win] = new(track)
		}
		t[win].Start = time.Now()
		prev = t[win]
	}
	return
}

func collect(tracks tracker) {
	X, err := xgb.NewConn()
	if err != nil {
		log.Fatal("xgb", err)
	}
	defer X.Close()

	err = screensaver.Init(X)
	if err != nil {
		log.Fatal("screensaver", err)
	}

	root := rootWin(X)

	drw := xproto.Drawable(root)
	screensaver.SelectInput(X, drw, screensaver.EventNotifyMask)

	spy(X, root)
	prev := tracks.Update(X, root)

	for {
		ev, everr := X.WaitForEvent()
		if everr != nil {
			log.Fatal("wait for event", everr)
		}
		switch event := ev.(type) {
		case xproto.PropertyNotifyEvent:
			if prev != nil {
				prev.Spent += time.Since(prev.Start)
			}
			prev = tracks.Update(X, root)
		case screensaver.NotifyEvent:
			switch event.State {
			case screensaver.StateOn:
				fmt.Println("away from keyboard")
				prev = nil
			}
		}
	}
}

func display(tracks tracker) {
	for {
		var total time.Duration
		classtotal := make(map[string]time.Duration)
		for n, t := range tracks {
			fmt.Println(n, t)
			total += t.Spent
			classtotal[n.Class] += t.Spent
		}
		fmt.Println("")
		for k, v := range classtotal {
			fmt.Println(k, v)
		}
		fmt.Println("Total:", total)
		fmt.Println("")
		time.Sleep(3 * time.Second)
	}
}

func cleanup(tracks tracker) {
	for {
		for k, v := range tracks {
			if time.Since(v.Start).Hours() > 12.0 {
				log.Println("removing", k)
				delete(tracks, k)
			}
		}
		time.Sleep(time.Minute)
	}
}

func main() {
	tracks := make(tracker)
	go collect(tracks)
	go cleanup(tracks)
	display(tracks)
}
<|endoftext|>


<commit_before>// Package convertkit provides a client to the ConvertKit API v3.
// See http://help.convertkit.com/article/33-api-documentation-v3
package convertkit

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"golang.org/x/sync/errgroup"
)

// Config is used to configure the creation of the client.
type Config struct {
	Endpoint string
	Key      string
	Secret   string

	HTTPClient *http.Client
}

// DefaultConfig returns a default configuration for the client. It parses the
// environment variables CONVERTKIT_ENDPOINT, CONVERTKIT_API_KEY, and
// CONVERTKIT_API_SECRET.
func DefaultConfig() *Config {
	c := Config{
		Endpoint:   "https://api.convertkit.com",
		HTTPClient: http.DefaultClient,
	}
	if v := os.Getenv("CONVERTKIT_ENDPOINT"); v != "" {
		c.Endpoint = v
	}
	if v := os.Getenv("CONVERTKIT_API_KEY"); v != "" {
		c.Key = v
	}
	if v := os.Getenv("CONVERTKIT_API_SECRET"); v != "" {
		c.Secret = v
	}
	return &c
}

// Client is the client to the ConvertKit API. Create a client with NewClient.
type Client struct {
	config *Config
}

// NewClient returns a new client for the given configuration.
func NewClient(c *Config) (*Client, error) {
	defConfig := DefaultConfig()
	if c.Endpoint == "" {
		c.Endpoint = defConfig.Endpoint
	}
	if c.Key == "" {
		c.Key = defConfig.Key
	}
	if c.Secret == "" {
		c.Secret = defConfig.Secret
	}
	if c.HTTPClient == nil {
		c.HTTPClient = defConfig.HTTPClient
	}
	return &Client{config: c}, nil
}

// Subscriber describes a ConvertKit subscriber.
type Subscriber struct {
	ID           int               `json:"id"`
	FirstName    string            `json:"first_name"`
	EmailAddress string            `json:"email_address"`
	State        string            `json:"state"`
	CreatedAt    time.Time         `json:"created_at"`
	Fields       map[string]string `json:"fields"`
}

type subscriberPage struct {
	TotalSubscribers int          `json:"total_subscribers"`
	Page             int          `json:"page"`
	TotalPages       int          `json:"total_pages"`
	Subscribers      []Subscriber `json:"subscribers"`
}

// Subscribers returns a list of all confirmed subscribers.
func (c *Client) Subscribers() ([]Subscriber, error) {
	p, err := c.subscriberPage(1)
	if err != nil {
		return nil, err
	}

	total := p.TotalPages
	pages := make([]subscriberPage, total)
	pages[0] = *p

	// TODO: limit number of Go routines to be nicer to the API
	var g errgroup.Group
	for i := 2; i <= total; i++ {
		i := i // see https://golang.org/doc/faq#closures_and_goroutines
		g.Go(func() error {
			p, err := c.subscriberPage(i)
			if err == nil {
				pages[i-1] = *p
			}
			return err
		})
	}

	if err := g.Wait(); err != nil {
		return nil, err
	}

	var subscribers []Subscriber
	for i := 0; i < total; i++ {
		subscribers = append(subscribers, pages[i].Subscribers...)
	}

	return subscribers, nil
}

// TotalSubscribers returns the number of confirmed subscribers.
func (c *Client) TotalSubscribers() (int, error) {
	resp, err := c.subscriberPage(1)
	if err != nil {
		return 0, err
	}
	return resp.TotalSubscribers, nil
}

func (c *Client) subscriberPage(page int) (*subscriberPage, error) {
	url := fmt.Sprintf("%s/v3/subscribers?api_secret=%s&page=%d",
		c.config.Endpoint, c.config.Secret, page)

	var resp subscriberPage
	if err := c.sendRequest("GET", url, nil, &resp); err != nil {
		return nil, err
	}

	return &resp, nil
}

func (c *Client) sendRequest(method, url string, body io.Reader, out interface{}) error {
	req, err := http.NewRequest(method, url, body)
	if err != nil {
		return err
	}

	resp, err := c.config.HTTPClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP error: %s", resp.Status)
	}

	return json.NewDecoder(resp.Body).Decode(out)
}
<commit_msg>Fix var names<commit_after>// Package convertkit provides a client to the ConvertKit API v3.
// See http://help.convertkit.com/article/33-api-documentation-v3
package convertkit

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"golang.org/x/sync/errgroup"
)

// Config is used to configure the creation of the client.
type Config struct {
	Endpoint string
	Key      string
	Secret   string

	HTTPClient *http.Client
}

// DefaultConfig returns a default configuration for the client. It parses the
// environment variables CONVERTKIT_ENDPOINT, CONVERTKIT_API_KEY, and
// CONVERTKIT_API_SECRET.
func DefaultConfig() *Config {
	c := Config{
		Endpoint:   "https://api.convertkit.com",
		HTTPClient: http.DefaultClient,
	}
	if v := os.Getenv("CONVERTKIT_ENDPOINT"); v != "" {
		c.Endpoint = v
	}
	if v := os.Getenv("CONVERTKIT_API_KEY"); v != "" {
		c.Key = v
	}
	if v := os.Getenv("CONVERTKIT_API_SECRET"); v != "" {
		c.Secret = v
	}
	return &c
}

// Client is the client to the ConvertKit API. Create a client with NewClient.
type Client struct {
	config *Config
}

// NewClient returns a new client for the given configuration.
func NewClient(c *Config) (*Client, error) {
	defConfig := DefaultConfig()
	if c.Endpoint == "" {
		c.Endpoint = defConfig.Endpoint
	}
	if c.Key == "" {
		c.Key = defConfig.Key
	}
	if c.Secret == "" {
		c.Secret = defConfig.Secret
	}
	if c.HTTPClient == nil {
		c.HTTPClient = defConfig.HTTPClient
	}
	return &Client{config: c}, nil
}

// Subscriber describes a ConvertKit subscriber.
type Subscriber struct {
	ID           int               `json:"id"`
	FirstName    string            `json:"first_name"`
	EmailAddress string            `json:"email_address"`
	State        string            `json:"state"`
	CreatedAt    time.Time         `json:"created_at"`
	Fields       map[string]string `json:"fields"`
}

type subscriberPage struct {
	TotalSubscribers int          `json:"total_subscribers"`
	Page             int          `json:"page"`
	TotalPages       int          `json:"total_pages"`
	Subscribers      []Subscriber `json:"subscribers"`
}

// Subscribers returns a list of all confirmed subscribers.
func (c *Client) Subscribers() ([]Subscriber, error) {
	p, err := c.subscriberPage(1)
	if err != nil {
		return nil, err
	}

	total := p.TotalPages
	pages := make([]subscriberPage, total)
	pages[0] = *p

	// TODO: limit number of Go routines to be nicer to the API
	var g errgroup.Group
	for i := 2; i <= total; i++ {
		i := i // see https://golang.org/doc/faq#closures_and_goroutines
		g.Go(func() error {
			p, err := c.subscriberPage(i)
			if err == nil {
				pages[i-1] = *p
			}
			return err
		})
	}

	if err := g.Wait(); err != nil {
		return nil, err
	}

	var subscribers []Subscriber
	for i := 0; i < total; i++ {
		subscribers = append(subscribers, pages[i].Subscribers...)
	}

	return subscribers, nil
}

// TotalSubscribers returns the number of confirmed subscribers.
func (c *Client) TotalSubscribers() (int, error) {
	p, err := c.subscriberPage(1)
	if err != nil {
		return 0, err
	}
	return p.TotalSubscribers, nil
}

func (c *Client) subscriberPage(page int) (*subscriberPage, error) {
	url := fmt.Sprintf("%s/v3/subscribers?api_secret=%s&page=%d",
		c.config.Endpoint, c.config.Secret, page)

	var p subscriberPage
	if err := c.sendRequest("GET", url, nil, &p); err != nil {
		return nil, err
	}

	return &p, nil
}

func (c *Client) sendRequest(method, url string, body io.Reader, out interface{}) error {
	req, err := http.NewRequest(method, url, body)
	if err != nil {
		return err
	}

	resp, err := c.config.HTTPClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP error: %s", resp.Status)
	}

	return json.NewDecoder(resp.Body).Decode(out)
}
<|endoftext|>


<commit_before>package barcode

import (
	"errors"
	"fmt"
	"image"
	"image/color"
	"math"
)

type wrapFunc func(x, y int) color.Color

type scaledBarcode struct {
	wrapped     Barcode
	wrapperFunc wrapFunc
	rect        image.Rectangle
}

func (bc *scaledBarcode) Content() string {
	return bc.wrapped.Content()
}

func (bc *scaledBarcode) Metadata() Metadata {
	return bc.wrapped.Metadata()
}

func (bc *scaledBarcode) ColorModel() color.Model {
	return bc.wrapped.ColorModel()
}

func (bc *scaledBarcode) Bounds() image.Rectangle {
	return bc.rect
}

func (bc *scaledBarcode) At(x, y int) color.Color {
	return bc.wrapperFunc(x, y)
}

// returns a resized barcode with the given width and height.
func Scale(bc Barcode, width, height int) (Barcode, error) {
	switch bc.Metadata().Dimensions {
	case 1:
		return scale1DCode(bc, width, height)
	case 2:
		return scale2DCode(bc, width, height)
	}

	return nil, errors.New("unsupported barcode format")
}

func scale2DCode(bc Barcode, width, height int) (Barcode, error) {
	orgBounds := bc.Bounds()
	orgWidth := orgBounds.Max.X - orgBounds.Min.X
	orgHeight := orgBounds.Max.Y - orgBounds.Min.Y

	factor := int(math.Min(float64(width)/float64(orgWidth), float64(height)/float64(orgHeight)))
	if factor <= 0 {
		return nil, fmt.Errorf("can no scale barcode to an image smaller then %dx%d", orgWidth, orgHeight)
	}

	offsetX := (width - (orgWidth * factor)) / 2
	offsetY := (height - (orgHeight * factor)) / 2

	wrap := func(x, y int) color.Color {
		if x < offsetX || y < offsetY {
			return color.White
		}
		x = (x - offsetX) / factor
		y = (y - offsetY) / factor
		if x >= orgWidth || y >= orgHeight {
			return color.White
		}
		return bc.At(x, y)
	}

	return &scaledBarcode{
		bc,
		wrap,
		image.Rect(0, 0, width, height),
	}, nil
}

func scale1DCode(bc Barcode, width, height int) (Barcode, error) {
	orgBounds := bc.Bounds()
	orgWidth := orgBounds.Max.X - orgBounds.Min.X
	factor := int(float64(width) / float64(orgWidth))

	if factor <= 0 {
		return nil, fmt.Errorf("can no scale barcode to an image smaller then %dx1", orgWidth)
	}
	offsetX := (width - (orgWidth * factor)) / 2

	wrap := func(x, y int) color.Color {
		if x < offsetX {
			return color.White
		}
		x = (x - offsetX) / factor

		if x >= orgWidth {
			return color.White
		}
		return bc.At(x, 0)
	}

	return &scaledBarcode{
		bc,
		wrap,
		image.Rect(0, 0, width, height),
	}, nil

}
<commit_msg>fixed typo<commit_after>package barcode

import (
	"errors"
	"fmt"
	"image"
	"image/color"
	"math"
)

type wrapFunc func(x, y int) color.Color

type scaledBarcode struct {
	wrapped     Barcode
	wrapperFunc wrapFunc
	rect        image.Rectangle
}

func (bc *scaledBarcode) Content() string {
	return bc.wrapped.Content()
}

func (bc *scaledBarcode) Metadata() Metadata {
	return bc.wrapped.Metadata()
}

func (bc *scaledBarcode) ColorModel() color.Model {
	return bc.wrapped.ColorModel()
}

func (bc *scaledBarcode) Bounds() image.Rectangle {
	return bc.rect
}

func (bc *scaledBarcode) At(x, y int) color.Color {
	return bc.wrapperFunc(x, y)
}

// returns a resized barcode with the given width and height.
func Scale(bc Barcode, width, height int) (Barcode, error) {
	switch bc.Metadata().Dimensions {
	case 1:
		return scale1DCode(bc, width, height)
	case 2:
		return scale2DCode(bc, width, height)
	}

	return nil, errors.New("unsupported barcode format")
}

func scale2DCode(bc Barcode, width, height int) (Barcode, error) {
	orgBounds := bc.Bounds()
	orgWidth := orgBounds.Max.X - orgBounds.Min.X
	orgHeight := orgBounds.Max.Y - orgBounds.Min.Y

	factor := int(math.Min(float64(width)/float64(orgWidth), float64(height)/float64(orgHeight)))
	if factor <= 0 {
		return nil, fmt.Errorf("can not scale barcode to an image smaller then %dx%d", orgWidth, orgHeight)
	}

	offsetX := (width - (orgWidth * factor)) / 2
	offsetY := (height - (orgHeight * factor)) / 2

	wrap := func(x, y int) color.Color {
		if x < offsetX || y < offsetY {
			return color.White
		}
		x = (x - offsetX) / factor
		y = (y - offsetY) / factor
		if x >= orgWidth || y >= orgHeight {
			return color.White
		}
		return bc.At(x, y)
	}

	return &scaledBarcode{
		bc,
		wrap,
		image.Rect(0, 0, width, height),
	}, nil
}

func scale1DCode(bc Barcode, width, height int) (Barcode, error) {
	orgBounds := bc.Bounds()
	orgWidth := orgBounds.Max.X - orgBounds.Min.X
	factor := int(float64(width) / float64(orgWidth))

	if factor <= 0 {
		return nil, fmt.Errorf("can not scale barcode to an image smaller then %dx1", orgWidth)
	}
	offsetX := (width - (orgWidth * factor)) / 2

	wrap := func(x, y int) color.Color {
		if x < offsetX {
			return color.White
		}
		x = (x - offsetX) / factor

		if x >= orgWidth {
			return color.White
		}
		return bc.At(x, 0)
	}

	return &scaledBarcode{
		bc,
		wrap,
		image.Rect(0, 0, width, height),
	}, nil

}
<|endoftext|>


<commit_before>package scans

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/ion-channel/ionic/vulnerabilities"
)

// UntranslatedResults represents a result of a specific type that has not been
// translated for use in reports
type UntranslatedResults struct {
	AboutYML                *AboutYMLResults                `json:"about_yml,omitempty"`
	Community               *CommunityResults               `json:"community,omitempty"`
	Coverage                *CoverageResults                `json:"coverage,omitempty"`
	Dependency              *DependencyResults              `json:"dependency,omitempty"`
	Difference              *DifferenceResults              `json:"difference,omitempty"`
	Ecosystem               *EcosystemResults               `json:"ecosystems,omitempty"`
	ExternalVulnerabilities *ExternalVulnerabilitiesResults `json:"external_vulnerability,omitempty"`
	License                 *LicenseResults                 `json:"license,omitempty"`
	Virus                   *VirusResults                   `json:"clamav,omitempty"`
	Vulnerability           *VulnerabilityResults           `json:"vulnerabilities,omitempty"`
}

// Translate moves information from the particular sub-struct, IE
// AboutYMLResults or LicenseResults into a generic, Data struct
func (u *UntranslatedResults) Translate() *TranslatedResults {
	var tr TranslatedResults
	// There is an argument to be made that the following "if" clauses
	// could be simplified with introspection since they all do
	// basically the same thing. I've (dmiles) chosen to writ it all
	// out in the name of explicit, easily-readable code.
	if u.AboutYML != nil {
		tr.Type = "about_yml"
		tr.Data = *u.AboutYML
	}
	if u.Community != nil {
		tr.Type = "community"
		tr.Data = *u.Community
	}
	if u.Coverage != nil {
		tr.Type = "coverage"
		tr.Data = *u.Coverage
	}
	if u.Dependency != nil {
		tr.Type = "dependency"
		tr.Data = *u.Dependency
	}
	if u.Difference != nil {
		tr.Type = "difference"
		tr.Data = *u.Difference
	}
	if u.Ecosystem != nil {
		tr.Type = "ecosystems"
		tr.Data = *u.Ecosystem
	}
	if u.ExternalVulnerabilities != nil {
		tr.Type = "external_vulnerability"
		tr.Data = *u.ExternalVulnerabilities
	}
	if u.License != nil {
		tr.Type = "license"
		tr.Data = *u.License
	}
	if u.Virus != nil {
		tr.Type = "virus"
		tr.Data = *u.Virus
	}
	if u.Vulnerability != nil {
		tr.Type = "vulnerability"
		tr.Data = *u.Vulnerability
	}
	return &tr
}

// TranslatedResults represents a result of a specific type that has been
// translated for use in reports
type TranslatedResults struct {
	Type string      `json:"type" xml:"type"`
	Data interface{} `json:"data,omitempty" xml:"data,omitempty"`
}

type translatedResults struct {
	Type    string          `json:"type"`
	RawData json.RawMessage `json:"data"`
}

// UnmarshalJSON is a custom JSON unmarshaller implementation for the standard
// go json package to know how to properly interpret ScanSummaryResults from
// JSON.
func (r *TranslatedResults) UnmarshalJSON(b []byte) error {
	var tr translatedResults
	err := json.Unmarshal(b, &tr)
	if err != nil {
		return err
	}

	r.Type = tr.Type

	switch strings.ToLower(tr.Type) {
	case "about_yml":
		var a AboutYMLResults
		err := json.Unmarshal(tr.RawData, &a)
		if err != nil {
			return fmt.Errorf("failed to unmarshall about yml results: %v", err)
		}

		r.Data = a
	case "community":
		var c CommunityResults
		err := json.Unmarshal(tr.RawData, &c)
		if err != nil {
			// Note: Could be a slice, needs to be fixed
			if strings.Contains(err.Error(), "cannot unmarshal array") {
				var sliceOfCommunityResults []CommunityResults
				err := json.Unmarshal(tr.RawData, &sliceOfCommunityResults)
				if err == nil {
					c = sliceOfCommunityResults[0]
					break
				}
			}
			return fmt.Errorf("failed to unmarshall community results: %v", err)
		}

		r.Data = c
	case "coverage", "external_coverage":
		var c CoverageResults
		err := json.Unmarshal(tr.RawData, &c)
		if err != nil {
			return fmt.Errorf("failed to unmarshall coverage results: %v", err)
		}

		r.Data = c
	case "dependency":
		var d DependencyResults
		err := json.Unmarshal(tr.RawData, &d)
		if err != nil {
			return fmt.Errorf("failed to unmarshall dependency results: %v", err)
		}

		r.Data = d
	case "ecosystems":
		var e EcosystemResults
		err := json.Unmarshal(tr.RawData, &e)
		if err != nil {
			return fmt.Errorf("failed to unmarshall ecosystems results: %v", err)
		}

		r.Data = e
	case "license":
		var l LicenseResults
		err := json.Unmarshal(tr.RawData, &l)
		if err != nil {
			return fmt.Errorf("failed to unmarshall license results: %v", err)
		}

		r.Data = l
	case "virus", "clamav":
		var v VirusResults
		err := json.Unmarshal(tr.RawData, &v)
		if err != nil {
			return fmt.Errorf("failed to unmarshall virus results: %v", err)
		}

		r.Data = v
	case "vulnerability":
		var v VulnerabilityResults
		err := json.Unmarshal(tr.RawData, &v)
		if err != nil {
			return fmt.Errorf("failed to unmarshall vulnerability results: %v", err)
		}

		r.Data = v
	case "external_vulnerability":
		var v ExternalVulnerabilitiesResults
		err := json.Unmarshal(tr.RawData, &v)
		if err != nil {
			return fmt.Errorf("failed to unmarshall external vulnerabilities results: %v", err)
		}

		r.Data = v
	case "difference":
		var v DifferenceResults
		err := json.Unmarshal(tr.RawData, &v)
		if err != nil {
			return fmt.Errorf("failed to unmarshall difference results: %v", err)
		}

		r.Data = v
	default:
		return fmt.Errorf("unsupported results type found: %v", tr.Type)
	}

	return nil
}

// UnmarshalJSON is a custom JSON unmarshaller implementation for the standard
// go json package to know how to properly interpret ScanSummaryResults from
// JSON.
func (u *UntranslatedResults) UnmarshalJSON(b []byte) error {
	// first look for results in the proper translated format
	// e.g. CommunityResults
	tr := &translatedResults{}
	err := json.Unmarshal(b, tr)
	if err != nil {
		// we have received invalid stringified json
		return fmt.Errorf("unable to unmarshal json")
	}

	// if there is a type and it is `community`
	// parse the data out
	if tr.Type == "community" {
		c := &CommunityResults{}
		err = json.Unmarshal(tr.RawData, c)
		if err != nil {
			return err
		}
		u.Community = c
		return nil
	}

	// it is not translated and not community
	// ur2 is required to keep the parser from
	// recursing here
	type ur2 UntranslatedResults
	err = json.Unmarshal(b, (*ur2)(u))
	if err != nil {
		// we have received invalid stringified json
		return fmt.Errorf("unable to unmarshal json")
	}

	return nil
}

// AboutYMLResults represents the data collected from the AboutYML scan.  It
// includes a message and whether or not the About YML file found was valid or
// not.
type AboutYMLResults struct {
	Message string `json:"message" xml:"message"`
	Valid   bool   `json:"valid" xml:"valid"`
	Content string `json:"content" xml:"content"`
}

// CommunityResults represents the data collected from a community scan.  It
// represents all known data regarding the open community of a software project
type CommunityResults struct {
	Committers int    `json:"committers" xml:"committers"`
	Name       string `json:"name" xml:"name"`
	URL        string `json:"url" xml:"url"`
}

// CoverageResults represents the data collected from a code coverage scan.  It
// includes the value of the code coverage seen for the project.
type CoverageResults struct {
	Value float64 `json:"value" xml:"value"`
}

// Dependency represents data for an individual requirement resolution
type Dependency struct {
	LatestVersion string `json:"latest_version" xml:"latest_version"`
	Org           string `json:"org" xml:"org"`
	Name          string `json:"name" xml:"name"`
	Type          string `json:"type" xml:"type"`
	Package       string `json:"package" xml:"package"`
	Version       string `json:"version" xml:"version"`
	Scope         string `json:"scope" xml:"scope"`
	Requirement   string `json:"requirement" xml:"requirement"`
}

// DependencyMeta represents data for a summary of all dependencies resolved
type DependencyMeta struct {
	FirstDegreeCount     int `json:"first_degree_count" xml:"first_degree_count"`
	NoVersionCount       int `json:"no_version_count" xml:"no_version_count"`
	TotalUniqueCount     int `json:"total_unique_count" xml:"total_unique_count"`
	UpdateAvailableCount int `json:"update_available_count" xml:"update_available_count"`
}

// DependencyResults represents the data collected from a dependency scan.  It
// includes a list of the dependencies seen and meta data counts about those
// dependencies seen.
type DependencyResults struct {
	Dependencies []Dependency   `json:"dependencies" xml:"dependencies"`
	Meta         DependencyMeta `json:"meta" xml:"meta"`
}

// DifferenceResults represents the checksum of a project.  It includes a checksum
// and flag indicating if there was a difference detected within that last 5 scans
type DifferenceResults struct {
	Checksum   string `json:"checksum" xml:"checksum"`
	Difference bool   `json:"difference" xml:"difference"`
}

// EcosystemResults represents the data collected from an ecosystems scan.  It
// include the name of the ecosystem and the number of lines seen for the given
// ecosystem.
type EcosystemResults struct {
	Ecosystems map[string]int `json:"ecosystems" xml:"ecosystems"`
}

// MarshalJSON meets the marshaller interface to custom wrangle an ecosystem
// result into the json shape
func (e EcosystemResults) MarshalJSON() ([]byte, error) {
	return json.Marshal(e.Ecosystems)
}

// UnmarshalJSON meets the unmarshaller interface to custom wrangle the
// ecosystem scan into an ecosystem result
func (e *EcosystemResults) UnmarshalJSON(b []byte) error {
	var m map[string]int
	err := json.Unmarshal(b, &m)
	if err != nil {
		return fmt.Errorf("failed to unmarshal ecosystem result: %v", err.Error())
	}

	e.Ecosystems = m
	return nil
}

// ExternalVulnerabilitiesResults represents the data collected from an external
// vulnerability scan.  It includes the number of each vulnerability criticality
// seen within the project.
type ExternalVulnerabilitiesResults struct {
	Critical int `json:"critical" xml:"critical"`
	High     int `json:"high" xml:"high"`
	Medium   int `json:"medium" xml:"medium"`
	Low      int `json:"low" xml:"low"`
}

// LicenseResults represents the data collected from a license scan.  It
// includes the name and type of each license seen within the project.
type LicenseResults struct {
	*License `json:"license" xml:"license"`
}

// License represents a name and slice of types of licenses seen in a given file
type License struct {
	Name string        `json:"name" xml:"name"`
	Type []LicenseType `json:"type" xml:"type"`
}

// LicenseType represents a type of license such as MIT, Apache 2.0, etc
type LicenseType struct {
	Name string `json:"name" xml:"name"`
}

// VirusResults represents the data collected from a virus scan.  It includes
// information of the viruses seen and the virus scanner used.
type VirusResults struct {
	KnownViruses       int    `json:"known_viruses" xml:"known_viruses"`
	EngineVersion      string `json:"engine_version" xml:"engine_version"`
	ScannedDirectories int    `json:"scanned_directories" xml:"scanned_directories"`
	ScannedFiles       int    `json:"scanned_files" xml:"scanned_files"`
	InfectedFiles      int    `json:"infected_files" xml:"infected_files"`
	DataScanned        string `json:"data_scanned" xml:"data_scanned"`
	DataRead           string `json:"data_read" xml:"data_read"`
	Time               string `json:"time" xml:"time"`
	FileNotes          struct {
		EmptyFiles []string `json:"empty_file" xml:"empty_file"`
	} `json:"file_notes" xml:"file_notes"`
	ClamavDetails struct {
		ClamavVersion   string `json:"clamav_version" xml:"clamav_version"`
		ClamavDbVersion string `json:"clamav_db_version" xml:"clamav_db_version"`
	} `json:"clam_av_details" xml:"clam_av_details"`
}

//VulnerabilityResults represents the data collected from a vulnerability scan.  It includes
// information of the vulnerabilities seen.
type VulnerabilityResults struct {
	Vulnerabilities []struct {
		ID              int                             `json:"id" xml:"id"`
		ExternalID      string                          `json:"external_id" xml:"external_id"`
		SourceID        int                             `json:"source_id" xml:"source_id"`
		Title           string                          `json:"title" xml:"title"`
		Name            string                          `json:"name" xml:"name"`
		Org             string                          `json:"org" xml:"org"`
		Version         string                          `json:"version" xml:"version"`
		Up              interface{}                     `json:"up" xml:"up"`
		Edition         interface{}                     `json:"edition" xml:"edition"`
		Aliases         []string                        `json:"aliases" xml:"aliases"`
		CreatedAt       time.Time                       `json:"created_at" xml:"created_at"`
		UpdatedAt       time.Time                       `json:"updated_at" xml:"updated_at"`
		References      interface{}                     `json:"references" xml:"references"`
		Part            interface{}                     `json:"part" xml:"part"`
		Language        interface{}                     `json:"language" xml:"language"`
		Vulnerabilities []vulnerabilities.Vulnerability `json:"vulnerabilities" xml:"vulnerabilities"`
	} `json:"vulnerabilities" xml:"vulnerabilities"`
	Meta struct {
		VulnerabilityCount int `json:"vulnerability_count" xml:"vulnerability_count"`
	} `json:"meta" xml:"meta"`
}
<commit_msg>Adding confidence to license data<commit_after>package scans

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/ion-channel/ionic/vulnerabilities"
)

// UntranslatedResults represents a result of a specific type that has not been
// translated for use in reports
type UntranslatedResults struct {
	AboutYML                *AboutYMLResults                `json:"about_yml,omitempty"`
	Community               *CommunityResults               `json:"community,omitempty"`
	Coverage                *CoverageResults                `json:"coverage,omitempty"`
	Dependency              *DependencyResults              `json:"dependency,omitempty"`
	Difference              *DifferenceResults              `json:"difference,omitempty"`
	Ecosystem               *EcosystemResults               `json:"ecosystems,omitempty"`
	ExternalVulnerabilities *ExternalVulnerabilitiesResults `json:"external_vulnerability,omitempty"`
	License                 *LicenseResults                 `json:"license,omitempty"`
	Virus                   *VirusResults                   `json:"clamav,omitempty"`
	Vulnerability           *VulnerabilityResults           `json:"vulnerabilities,omitempty"`
}

// Translate moves information from the particular sub-struct, IE
// AboutYMLResults or LicenseResults into a generic, Data struct
func (u *UntranslatedResults) Translate() *TranslatedResults {
	var tr TranslatedResults
	// There is an argument to be made that the following "if" clauses
	// could be simplified with introspection since they all do
	// basically the same thing. I've (dmiles) chosen to writ it all
	// out in the name of explicit, easily-readable code.
	if u.AboutYML != nil {
		tr.Type = "about_yml"
		tr.Data = *u.AboutYML
	}
	if u.Community != nil {
		tr.Type = "community"
		tr.Data = *u.Community
	}
	if u.Coverage != nil {
		tr.Type = "coverage"
		tr.Data = *u.Coverage
	}
	if u.Dependency != nil {
		tr.Type = "dependency"
		tr.Data = *u.Dependency
	}
	if u.Difference != nil {
		tr.Type = "difference"
		tr.Data = *u.Difference
	}
	if u.Ecosystem != nil {
		tr.Type = "ecosystems"
		tr.Data = *u.Ecosystem
	}
	if u.ExternalVulnerabilities != nil {
		tr.Type = "external_vulnerability"
		tr.Data = *u.ExternalVulnerabilities
	}
	if u.License != nil {
		tr.Type = "license"
		tr.Data = *u.License
	}
	if u.Virus != nil {
		tr.Type = "virus"
		tr.Data = *u.Virus
	}
	if u.Vulnerability != nil {
		tr.Type = "vulnerability"
		tr.Data = *u.Vulnerability
	}
	return &tr
}

// TranslatedResults represents a result of a specific type that has been
// translated for use in reports
type TranslatedResults struct {
	Type string      `json:"type" xml:"type"`
	Data interface{} `json:"data,omitempty" xml:"data,omitempty"`
}

type translatedResults struct {
	Type    string          `json:"type"`
	RawData json.RawMessage `json:"data"`
}

// UnmarshalJSON is a custom JSON unmarshaller implementation for the standard
// go json package to know how to properly interpret ScanSummaryResults from
// JSON.
func (r *TranslatedResults) UnmarshalJSON(b []byte) error {
	var tr translatedResults
	err := json.Unmarshal(b, &tr)
	if err != nil {
		return err
	}

	r.Type = tr.Type

	switch strings.ToLower(tr.Type) {
	case "about_yml":
		var a AboutYMLResults
		err := json.Unmarshal(tr.RawData, &a)
		if err != nil {
			return fmt.Errorf("failed to unmarshall about yml results: %v", err)
		}

		r.Data = a
	case "community":
		var c CommunityResults
		err := json.Unmarshal(tr.RawData, &c)
		if err != nil {
			// Note: Could be a slice, needs to be fixed
			if strings.Contains(err.Error(), "cannot unmarshal array") {
				var sliceOfCommunityResults []CommunityResults
				err := json.Unmarshal(tr.RawData, &sliceOfCommunityResults)
				if err == nil {
					c = sliceOfCommunityResults[0]
					break
				}
			}
			return fmt.Errorf("failed to unmarshall community results: %v", err)
		}

		r.Data = c
	case "coverage", "external_coverage":
		var c CoverageResults
		err := json.Unmarshal(tr.RawData, &c)
		if err != nil {
			return fmt.Errorf("failed to unmarshall coverage results: %v", err)
		}

		r.Data = c
	case "dependency":
		var d DependencyResults
		err := json.Unmarshal(tr.RawData, &d)
		if err != nil {
			return fmt.Errorf("failed to unmarshall dependency results: %v", err)
		}

		r.Data = d
	case "ecosystems":
		var e EcosystemResults
		err := json.Unmarshal(tr.RawData, &e)
		if err != nil {
			return fmt.Errorf("failed to unmarshall ecosystems results: %v", err)
		}

		r.Data = e
	case "license":
		var l LicenseResults
		err := json.Unmarshal(tr.RawData, &l)
		if err != nil {
			return fmt.Errorf("failed to unmarshall license results: %v", err)
		}

		r.Data = l
	case "virus", "clamav":
		var v VirusResults
		err := json.Unmarshal(tr.RawData, &v)
		if err != nil {
			return fmt.Errorf("failed to unmarshall virus results: %v", err)
		}

		r.Data = v
	case "vulnerability":
		var v VulnerabilityResults
		err := json.Unmarshal(tr.RawData, &v)
		if err != nil {
			return fmt.Errorf("failed to unmarshall vulnerability results: %v", err)
		}

		r.Data = v
	case "external_vulnerability":
		var v ExternalVulnerabilitiesResults
		err := json.Unmarshal(tr.RawData, &v)
		if err != nil {
			return fmt.Errorf("failed to unmarshall external vulnerabilities results: %v", err)
		}

		r.Data = v
	case "difference":
		var v DifferenceResults
		err := json.Unmarshal(tr.RawData, &v)
		if err != nil {
			return fmt.Errorf("failed to unmarshall difference results: %v", err)
		}

		r.Data = v
	default:
		return fmt.Errorf("unsupported results type found: %v", tr.Type)
	}

	return nil
}

// UnmarshalJSON is a custom JSON unmarshaller implementation for the standard
// go json package to know how to properly interpret ScanSummaryResults from
// JSON.
func (u *UntranslatedResults) UnmarshalJSON(b []byte) error {
	// first look for results in the proper translated format
	// e.g. CommunityResults
	tr := &translatedResults{}
	err := json.Unmarshal(b, tr)
	if err != nil {
		// we have received invalid stringified json
		return fmt.Errorf("unable to unmarshal json")
	}

	// if there is a type and it is `community`
	// parse the data out
	if tr.Type == "community" {
		c := &CommunityResults{}
		err = json.Unmarshal(tr.RawData, c)
		if err != nil {
			return err
		}
		u.Community = c
		return nil
	}

	// it is not translated and not community
	// ur2 is required to keep the parser from
	// recursing here
	type ur2 UntranslatedResults
	err = json.Unmarshal(b, (*ur2)(u))
	if err != nil {
		// we have received invalid stringified json
		return fmt.Errorf("unable to unmarshal json")
	}

	return nil
}

// AboutYMLResults represents the data collected from the AboutYML scan.  It
// includes a message and whether or not the About YML file found was valid or
// not.
type AboutYMLResults struct {
	Message string `json:"message" xml:"message"`
	Valid   bool   `json:"valid" xml:"valid"`
	Content string `json:"content" xml:"content"`
}

// CommunityResults represents the data collected from a community scan.  It
// represents all known data regarding the open community of a software project
type CommunityResults struct {
	Committers int    `json:"committers" xml:"committers"`
	Name       string `json:"name" xml:"name"`
	URL        string `json:"url" xml:"url"`
}

// CoverageResults represents the data collected from a code coverage scan.  It
// includes the value of the code coverage seen for the project.
type CoverageResults struct {
	Value float64 `json:"value" xml:"value"`
}

// Dependency represents data for an individual requirement resolution
type Dependency struct {
	LatestVersion string `json:"latest_version" xml:"latest_version"`
	Org           string `json:"org" xml:"org"`
	Name          string `json:"name" xml:"name"`
	Type          string `json:"type" xml:"type"`
	Package       string `json:"package" xml:"package"`
	Version       string `json:"version" xml:"version"`
	Scope         string `json:"scope" xml:"scope"`
	Requirement   string `json:"requirement" xml:"requirement"`
}

// DependencyMeta represents data for a summary of all dependencies resolved
type DependencyMeta struct {
	FirstDegreeCount     int `json:"first_degree_count" xml:"first_degree_count"`
	NoVersionCount       int `json:"no_version_count" xml:"no_version_count"`
	TotalUniqueCount     int `json:"total_unique_count" xml:"total_unique_count"`
	UpdateAvailableCount int `json:"update_available_count" xml:"update_available_count"`
}

// DependencyResults represents the data collected from a dependency scan.  It
// includes a list of the dependencies seen and meta data counts about those
// dependencies seen.
type DependencyResults struct {
	Dependencies []Dependency   `json:"dependencies" xml:"dependencies"`
	Meta         DependencyMeta `json:"meta" xml:"meta"`
}

// DifferenceResults represents the checksum of a project.  It includes a checksum
// and flag indicating if there was a difference detected within that last 5 scans
type DifferenceResults struct {
	Checksum   string `json:"checksum" xml:"checksum"`
	Difference bool   `json:"difference" xml:"difference"`
}

// EcosystemResults represents the data collected from an ecosystems scan.  It
// include the name of the ecosystem and the number of lines seen for the given
// ecosystem.
type EcosystemResults struct {
	Ecosystems map[string]int `json:"ecosystems" xml:"ecosystems"`
}

// MarshalJSON meets the marshaller interface to custom wrangle an ecosystem
// result into the json shape
func (e EcosystemResults) MarshalJSON() ([]byte, error) {
	return json.Marshal(e.Ecosystems)
}

// UnmarshalJSON meets the unmarshaller interface to custom wrangle the
// ecosystem scan into an ecosystem result
func (e *EcosystemResults) UnmarshalJSON(b []byte) error {
	var m map[string]int
	err := json.Unmarshal(b, &m)
	if err != nil {
		return fmt.Errorf("failed to unmarshal ecosystem result: %v", err.Error())
	}

	e.Ecosystems = m
	return nil
}

// ExternalVulnerabilitiesResults represents the data collected from an external
// vulnerability scan.  It includes the number of each vulnerability criticality
// seen within the project.
type ExternalVulnerabilitiesResults struct {
	Critical int `json:"critical" xml:"critical"`
	High     int `json:"high" xml:"high"`
	Medium   int `json:"medium" xml:"medium"`
	Low      int `json:"low" xml:"low"`
}

// LicenseResults represents the data collected from a license scan.  It
// includes the name and type of each license seen within the project.
type LicenseResults struct {
	*License `json:"license" xml:"license"`
}

// License represents a name and slice of types of licenses seen in a given file
type License struct {
	Name string        `json:"name" xml:"name"`
	Type []LicenseType `json:"type" xml:"type"`
}

// LicenseType represents a type of license such as MIT, Apache 2.0, etc
type LicenseType struct {
	Name       string  `json:"name" xml:"name"`
	Confidence float32 `json:"confidence"`
}

// VirusResults represents the data collected from a virus scan.  It includes
// information of the viruses seen and the virus scanner used.
type VirusResults struct {
	KnownViruses       int    `json:"known_viruses" xml:"known_viruses"`
	EngineVersion      string `json:"engine_version" xml:"engine_version"`
	ScannedDirectories int    `json:"scanned_directories" xml:"scanned_directories"`
	ScannedFiles       int    `json:"scanned_files" xml:"scanned_files"`
	InfectedFiles      int    `json:"infected_files" xml:"infected_files"`
	DataScanned        string `json:"data_scanned" xml:"data_scanned"`
	DataRead           string `json:"data_read" xml:"data_read"`
	Time               string `json:"time" xml:"time"`
	FileNotes          struct {
		EmptyFiles []string `json:"empty_file" xml:"empty_file"`
	} `json:"file_notes" xml:"file_notes"`
	ClamavDetails struct {
		ClamavVersion   string `json:"clamav_version" xml:"clamav_version"`
		ClamavDbVersion string `json:"clamav_db_version" xml:"clamav_db_version"`
	} `json:"clam_av_details" xml:"clam_av_details"`
}

//VulnerabilityResults represents the data collected from a vulnerability scan.  It includes
// information of the vulnerabilities seen.
type VulnerabilityResults struct {
	Vulnerabilities []struct {
		ID              int                             `json:"id" xml:"id"`
		ExternalID      string                          `json:"external_id" xml:"external_id"`
		SourceID        int                             `json:"source_id" xml:"source_id"`
		Title           string                          `json:"title" xml:"title"`
		Name            string                          `json:"name" xml:"name"`
		Org             string                          `json:"org" xml:"org"`
		Version         string                          `json:"version" xml:"version"`
		Up              interface{}                     `json:"up" xml:"up"`
		Edition         interface{}                     `json:"edition" xml:"edition"`
		Aliases         []string                        `json:"aliases" xml:"aliases"`
		CreatedAt       time.Time                       `json:"created_at" xml:"created_at"`
		UpdatedAt       time.Time                       `json:"updated_at" xml:"updated_at"`
		References      interface{}                     `json:"references" xml:"references"`
		Part            interface{}                     `json:"part" xml:"part"`
		Language        interface{}                     `json:"language" xml:"language"`
		Vulnerabilities []vulnerabilities.Vulnerability `json:"vulnerabilities" xml:"vulnerabilities"`
	} `json:"vulnerabilities" xml:"vulnerabilities"`
	Meta struct {
		VulnerabilityCount int `json:"vulnerability_count" xml:"vulnerability_count"`
	} `json:"meta" xml:"meta"`
}
<|endoftext|>