<commit_before>package dnssd

import (
	"context"
	"fmt"
	"net"

	"github.com/brutella/dnssd/log"
	"github.com/miekg/dns"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

var (
	IPv4LinkLocalMulticast = net.ParseIP("224.0.0.251")
	IPv6LinkLocalMulticast = net.ParseIP("ff02::fb")

	AddrIPv4LinkLocalMulticast = &net.UDPAddr{
		IP:   IPv4LinkLocalMulticast,
		Port: 5353,
	}

	AddrIPv6LinkLocalMulticast = &net.UDPAddr{
		IP:   IPv6LinkLocalMulticast,
		Port: 5353,
	}

	TTLDefault  uint32 = 75 * 60 // Default ttl for mDNS resource records
	TTLHostname uint32 = 120     // TTL for mDNS resource records containing the host name
)

// Query is a mDNS query
type Query struct {
	msg   *dns.Msg       // The query message
	iface *net.Interface // The network interface to which the message is sent
}

// Response is a mDNS response
type Response struct {
	msg   *dns.Msg       // The response message
	addr  *net.UDPAddr   // Is nil for multicast response
	iface *net.Interface // The network interface to which the message is sent
}

// Request represents an incoming mDNS message
type Request struct {
	msg   *dns.Msg       // The message
	from  *net.UDPAddr   // The source addr of the message
	iface *net.Interface // The network interface from which the message was received
}

func (r Request) String() string {
	return fmt.Sprintf("%s@%s\n%v", r.from.IP, r.iface.Name, r.msg)
}

// MDNSConn represents a mDNS connection. It encapsulates an IPv4 and IPv6 UDP connection.
type MDNSConn interface {
	// SendQuery sends a mDNS query.
	SendQuery(q *Query) error

	// SendResponse sends a mDNS response
	SendResponse(resp *Response) error

	// Read returns a channel which receives mDNS messages
	Read(ctx context.Context) <-chan *Request

	// Clears the connection buffer
	Drain(ctx context.Context)

	// Close closes the connection
	Close()
}

type mdnsConn struct {
	ipv4 *ipv4.PacketConn
	ipv6 *ipv6.PacketConn
	ch   chan *Request
}

func NewMDNSConn() (MDNSConn, error) {
	return newMDNSConn()
}

func (c *mdnsConn) SendQuery(q *Query) error {
	return c.sendQuery(q.msg, q.iface)
}

func (c *mdnsConn) SendResponse(resp *Response) error {
	if resp.addr != nil {
		return c.sendResponseTo(resp.msg, resp.iface, resp.addr)
	}

	return c.sendResponse(resp.msg, resp.iface)
}

func (c *mdnsConn) Read(ctx context.Context) <-chan *Request {
	return c.read(ctx)
}

func (c *mdnsConn) Drain(ctx context.Context) {
	log.Debug.Println("Draining connection")
	for {
		select {
		case req := <-c.Read(ctx):
			log.Debug.Println("Ignoring msg from", req.from.IP)
			break
		default:
			return
		}
	}
}

func (c *mdnsConn) Close() {
	c.close()
}

func newMDNSConn() (*mdnsConn, error) {
	var errs []error
	var connIPv4 *ipv4.PacketConn
	var connIPv6 *ipv6.PacketConn

	if conn, err := net.ListenUDP("udp4", AddrIPv4LinkLocalMulticast); err != nil {
		errs = append(errs, err)
	} else {
		connIPv4 = ipv4.NewPacketConn(conn)
		connIPv4.SetControlMessage(ipv4.FlagInterface, true)
		connIPv4.SetMulticastLoopback(false)

		for _, iface := range multicastInterfaces() {
			if err := connIPv4.JoinGroup(&iface, &net.UDPAddr{IP: IPv4LinkLocalMulticast}); err != nil {
				log.Debug.Printf("Failed joining IPv4 %v: %v", iface.Name, err)
			} else {
				log.Debug.Printf("Joined IPv4 %v", iface.Name)
			}
		}
	}

	if conn, err := net.ListenUDP("udp6", AddrIPv6LinkLocalMulticast); err != nil {
		errs = append(errs, err)
	} else {
		connIPv6 = ipv6.NewPacketConn(conn)
		connIPv6.SetControlMessage(ipv6.FlagInterface, true)
		connIPv6.SetMulticastLoopback(false)

		for _, iface := range multicastInterfaces() {
			if err := connIPv6.JoinGroup(&iface, &net.UDPAddr{IP: IPv6LinkLocalMulticast}); err != nil {
				log.Debug.Printf("Failed joining IPv6 %v: %v", iface.Name, err)
			} else {
				log.Debug.Printf("Joined IPv6 %v", iface.Name)
			}
		}
	}

	if err := first(errs...); connIPv4 == nil && connIPv6 == nil {
		return nil, fmt.Errorf("Failed setting up UDP server: %v", err)
	}

	return &mdnsConn{
		ipv4: connIPv4,
		ipv6: connIPv6,
		ch:   make(chan *Request),
	}, nil
}

func (c *mdnsConn) enableMulticastLoopback(enable bool) {
	if c.ipv4 != nil {
		c.ipv4.SetMulticastLoopback(enable)
	}

	if c.ipv6 != nil {
		c.ipv6.SetMulticastLoopback(enable)
	}
}

func (c *mdnsConn) close() {
	if c.ipv4 != nil {
		c.ipv4.Close()
	}

	if c.ipv6 != nil {
		c.ipv6.Close()
	}
}

func (c *mdnsConn) read(ctx context.Context) <-chan *Request {
	c.readInto(ctx, c.ch)
	return c.ch
}

func (c *mdnsConn) readInto(ctx context.Context, ch chan *Request) {

	isDone := func(ctx context.Context) bool {
		return ctx.Err() == context.Canceled
	}

	if c.ipv4 != nil {
		go func() {
			buf := make([]byte, 65536)
			for {
				if isDone(ctx) {
					return
				}

				n, cm, from, err := c.ipv4.ReadFrom(buf)
				if err != nil {
					continue
				}

				udpAddr, ok := from.(*net.UDPAddr)
				if !ok {
					log.Info.Println("invalid source address")
					continue
				}

				var iface *net.Interface
				if cm != nil {
					iface, err = net.InterfaceByIndex(cm.IfIndex)
					if err != nil {
						continue
					}
				}

				if n > 0 {
					m := new(dns.Msg)
					if err := m.Unpack(buf); err == nil && !shouldIgnore(m) {
						ch <- &Request{m, udpAddr, iface}
					}
				}
			}
		}()
	}

	if c.ipv6 != nil {
		go func() {
			buf := make([]byte, 65536)
			for {
				if isDone(ctx) {
					return
				}

				n, cm, from, err := c.ipv6.ReadFrom(buf)
				if err != nil {
					continue
				}

				udpAddr, ok := from.(*net.UDPAddr)
				if !ok {
					log.Info.Println("invalid source address")
					continue
				}

				var iface *net.Interface
				if cm != nil {
					iface, err = net.InterfaceByIndex(cm.IfIndex)
					if err != nil {
						continue
					}
				}

				if n > 0 {
					m := new(dns.Msg)
					if err := m.Unpack(buf); err == nil && !shouldIgnore(m) {
						ch <- &Request{m, udpAddr, iface}
					}
				}
			}
		}()
	}
}

func (c *mdnsConn) sendQuery(m *dns.Msg, iface *net.Interface) error {
	sanitizeQuery(m)

	return c.writeMsg(m, iface)
}

func (c *mdnsConn) sendResponse(m *dns.Msg, iface *net.Interface) error {
	sanitizeResponse(m)

	return c.writeMsg(m, iface)
}

func (c *mdnsConn) sendResponseTo(m *dns.Msg, iface *net.Interface, addr *net.UDPAddr) error {
	sanitizeResponse(m)

	return c.writeMsgTo(m, iface, addr)
}

func (c *mdnsConn) writeMsg(m *dns.Msg, iface *net.Interface) error {
	var err error
	if c.ipv4 != nil {
		err = c.writeMsgTo(m, iface, AddrIPv4LinkLocalMulticast)
	}

	if c.ipv6 != nil {
		err = c.writeMsgTo(m, iface, AddrIPv6LinkLocalMulticast)
	}

	return err
}

func (c *mdnsConn) writeMsgTo(m *dns.Msg, iface *net.Interface, addr *net.UDPAddr) error {
	sanitizeMsg(m)

	var err error
	if c.ipv4 != nil && addr.IP.To4() != nil {
		if out, err := m.Pack(); err == nil {
			ctrl := &ipv4.ControlMessage{}
			if iface != nil {
				ctrl.IfIndex = iface.Index
			}
			_, err = c.ipv4.WriteTo(out, ctrl, addr)
		}
	}

	if c.ipv6 != nil && addr.IP.To4() == nil {
		if out, err := m.Pack(); err == nil {
			ctrl := &ipv6.ControlMessage{}
			if iface != nil {
				ctrl.IfIndex = iface.Index
			}
			_, err = c.ipv6.WriteTo(out, ctrl, addr)
		}
	}

	return err
}

func shouldIgnore(m *dns.Msg) bool {
	if m.Opcode != 0 {
		return true
	}

	if m.Rcode != 0 {
		return true
	}

	return false
}

func sanitizeResponse(m *dns.Msg) {
	if m.Question != nil && len(m.Question) > 0 {
		log.Info.Println("Multicast DNS responses MUST NOT contain any questions in the Question Section.  (RFC6762 6)")
		m.Question = nil
	}

	if !m.Response {
		log.Info.Println("In response messages the QR bit MUST be one (RFC6762 18.2)")
		m.Response = true
	}

	if !m.Authoritative {
		log.Info.Println("AA Bit bit MUST be set to one in response messages (RFC6762 18.4)")
		m.Authoritative = true
	}

	if m.Truncated {
		log.Info.Println("In multicast response messages, the TC bit MUST be zero on transmission. (RFC6762 18.5)")
		m.Truncated = false
	}
}

func sanitizeQuery(m *dns.Msg) {
	if m.Response {
		log.Info.Println("In query messages the QR bit MUST be zero (RFC6762 18.2)")
		m.Response = false
	}

	if m.Authoritative {
		log.Info.Println("AA Bit MUST be zero in query messages (RFC6762 18.4)")
		m.Authoritative = false
	}
}

func sanitizeMsg(m *dns.Msg) {
	if m.Opcode != 0 {
		log.Info.Println("In both multicast query and multicast response messages, the OPCODE MUST be zero on transmission (RFC6762 18.3)")
		m.Opcode = 0
	}

	if m.RecursionDesired {
		log.Info.Println("In both multicast query and multicast response messages, the Recursion Available bit MUST be zero on transmission. (RFC6762 18.7)")
		m.RecursionDesired = false
	}

	if m.Zero {
		log.Info.Println("In both query and response messages, the Zero bit MUST be zero on transmission (RFC6762 18.8)")
		m.Zero = false
	}

	if m.AuthenticatedData {
		log.Info.Println("In both multicast query and multicast response messages, the Authentic Data bit MUST be zero on transmission (RFC6762 18.9)")
		m.AuthenticatedData = false
	}

	if m.CheckingDisabled {
		log.Info.Println("In both multicast query and multicast response messages, the Checking Disabled bit MUST be zero on transmission (RFC6762 18.10)")
		m.CheckingDisabled = false
	}

	if m.Rcode != 0 {
		log.Info.Println("In both multicast query and multicast response messages, the Response Code MUST be zero on transmission. (RFC6762 18.11)")
		m.Rcode = 0
	}
}

func first(errs ...error) error {
	for _, err := range errs {
		if err != nil {
			return err
		}
	}

	return nil
}

// Sets the Top Bit of rrclass for all answer records (except PTR) to trigger a cache flush in the receivers.
func setAnswerCacheFlushBit(msg *dns.Msg) {
	// From RFC6762
	//    The most significant bit of the rrclass for a record in the Answer
	//    Section of a response message is the Multicast DNS cache-flush bit
	//    and is discussed in more detail below in Section 10.2, "Announcements
	//    to Flush Outdated Cache Entries".
	for _, a := range msg.Answer {
		switch a.(type) {
		case *dns.PTR:
			continue
		default:
			a.Header().Class |= (1 << 15)
		}
	}
}

// Sets the Top Bit of class to indicate the unicast responses are preferred for this question.
func setQuestionUnicast(q *dns.Question) {
	q.Qclass |= (1 << 15)
}

// Returns true if q requires unicast responses.
func isUnicastQuestion(q dns.Question) bool {
	// From RFC6762
	// 18.12.  Repurposing of Top Bit of qclass in Question Section
	//
	//    In the Question Section of a Multicast DNS query, the top bit of the
	//    qclass field is used to indicate that unicast responses are preferred
	//    for this particular question.  (See Section 5.4.)
	return q.Qclass&(1<<15) != 0
}
<commit_msg>Use no control message when no interface specified<commit_after>package dnssd

import (
	"context"
	"fmt"
	"net"

	"github.com/brutella/dnssd/log"
	"github.com/miekg/dns"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

var (
	IPv4LinkLocalMulticast = net.ParseIP("224.0.0.251")
	IPv6LinkLocalMulticast = net.ParseIP("ff02::fb")

	AddrIPv4LinkLocalMulticast = &net.UDPAddr{
		IP:   IPv4LinkLocalMulticast,
		Port: 5353,
	}

	AddrIPv6LinkLocalMulticast = &net.UDPAddr{
		IP:   IPv6LinkLocalMulticast,
		Port: 5353,
	}

	TTLDefault  uint32 = 75 * 60 // Default ttl for mDNS resource records
	TTLHostname uint32 = 120     // TTL for mDNS resource records containing the host name
)

// Query is a mDNS query
type Query struct {
	msg   *dns.Msg       // The query message
	iface *net.Interface // The network interface to which the message is sent
}

// Response is a mDNS response
type Response struct {
	msg   *dns.Msg       // The response message
	addr  *net.UDPAddr   // Is nil for multicast response
	iface *net.Interface // The network interface to which the message is sent
}

// Request represents an incoming mDNS message
type Request struct {
	msg   *dns.Msg       // The message
	from  *net.UDPAddr   // The source addr of the message
	iface *net.Interface // The network interface from which the message was received
}

func (r Request) String() string {
	return fmt.Sprintf("%s@%s\n%v", r.from.IP, r.iface.Name, r.msg)
}

// MDNSConn represents a mDNS connection. It encapsulates an IPv4 and IPv6 UDP connection.
type MDNSConn interface {
	// SendQuery sends a mDNS query.
	SendQuery(q *Query) error

	// SendResponse sends a mDNS response
	SendResponse(resp *Response) error

	// Read returns a channel which receives mDNS messages
	Read(ctx context.Context) <-chan *Request

	// Clears the connection buffer
	Drain(ctx context.Context)

	// Close closes the connection
	Close()
}

type mdnsConn struct {
	ipv4 *ipv4.PacketConn
	ipv6 *ipv6.PacketConn
	ch   chan *Request
}

func NewMDNSConn() (MDNSConn, error) {
	return newMDNSConn()
}

func (c *mdnsConn) SendQuery(q *Query) error {
	return c.sendQuery(q.msg, q.iface)
}

func (c *mdnsConn) SendResponse(resp *Response) error {
	if resp.addr != nil {
		return c.sendResponseTo(resp.msg, resp.iface, resp.addr)
	}

	return c.sendResponse(resp.msg, resp.iface)
}

func (c *mdnsConn) Read(ctx context.Context) <-chan *Request {
	return c.read(ctx)
}

func (c *mdnsConn) Drain(ctx context.Context) {
	log.Debug.Println("Draining connection")
	for {
		select {
		case req := <-c.Read(ctx):
			log.Debug.Println("Ignoring msg from", req.from.IP)
			break
		default:
			return
		}
	}
}

func (c *mdnsConn) Close() {
	c.close()
}

func newMDNSConn() (*mdnsConn, error) {
	var errs []error
	var connIPv4 *ipv4.PacketConn
	var connIPv6 *ipv6.PacketConn

	if conn, err := net.ListenUDP("udp4", AddrIPv4LinkLocalMulticast); err != nil {
		errs = append(errs, err)
	} else {
		connIPv4 = ipv4.NewPacketConn(conn)
		connIPv4.SetControlMessage(ipv4.FlagInterface, true)
		connIPv4.SetMulticastLoopback(false)

		for _, iface := range multicastInterfaces() {
			if err := connIPv4.JoinGroup(&iface, &net.UDPAddr{IP: IPv4LinkLocalMulticast}); err != nil {
				log.Debug.Printf("Failed joining IPv4 %v: %v", iface.Name, err)
			} else {
				log.Debug.Printf("Joined IPv4 %v", iface.Name)
			}
		}
	}

	if conn, err := net.ListenUDP("udp6", AddrIPv6LinkLocalMulticast); err != nil {
		errs = append(errs, err)
	} else {
		connIPv6 = ipv6.NewPacketConn(conn)
		connIPv6.SetControlMessage(ipv6.FlagInterface, true)
		connIPv6.SetMulticastLoopback(false)

		for _, iface := range multicastInterfaces() {
			if err := connIPv6.JoinGroup(&iface, &net.UDPAddr{IP: IPv6LinkLocalMulticast}); err != nil {
				log.Debug.Printf("Failed joining IPv6 %v: %v", iface.Name, err)
			} else {
				log.Debug.Printf("Joined IPv6 %v", iface.Name)
			}
		}
	}

	if err := first(errs...); connIPv4 == nil && connIPv6 == nil {
		return nil, fmt.Errorf("Failed setting up UDP server: %v", err)
	}

	return &mdnsConn{
		ipv4: connIPv4,
		ipv6: connIPv6,
		ch:   make(chan *Request),
	}, nil
}

func (c *mdnsConn) enableMulticastLoopback(enable bool) {
	if c.ipv4 != nil {
		c.ipv4.SetMulticastLoopback(enable)
	}

	if c.ipv6 != nil {
		c.ipv6.SetMulticastLoopback(enable)
	}
}

func (c *mdnsConn) close() {
	if c.ipv4 != nil {
		c.ipv4.Close()
	}

	if c.ipv6 != nil {
		c.ipv6.Close()
	}
}

func (c *mdnsConn) read(ctx context.Context) <-chan *Request {
	c.readInto(ctx, c.ch)
	return c.ch
}

func (c *mdnsConn) readInto(ctx context.Context, ch chan *Request) {

	isDone := func(ctx context.Context) bool {
		return ctx.Err() == context.Canceled
	}

	if c.ipv4 != nil {
		go func() {
			buf := make([]byte, 65536)
			for {
				if isDone(ctx) {
					return
				}

				n, cm, from, err := c.ipv4.ReadFrom(buf)
				if err != nil {
					continue
				}

				udpAddr, ok := from.(*net.UDPAddr)
				if !ok {
					log.Info.Println("invalid source address")
					continue
				}

				var iface *net.Interface
				if cm != nil {
					iface, err = net.InterfaceByIndex(cm.IfIndex)
					if err != nil {
						continue
					}
				}

				if n > 0 {
					m := new(dns.Msg)
					if err := m.Unpack(buf); err == nil && !shouldIgnore(m) {
						ch <- &Request{m, udpAddr, iface}
					}
				}
			}
		}()
	}

	if c.ipv6 != nil {
		go func() {
			buf := make([]byte, 65536)
			for {
				if isDone(ctx) {
					return
				}

				n, cm, from, err := c.ipv6.ReadFrom(buf)
				if err != nil {
					continue
				}

				udpAddr, ok := from.(*net.UDPAddr)
				if !ok {
					log.Info.Println("invalid source address")
					continue
				}

				var iface *net.Interface
				if cm != nil {
					iface, err = net.InterfaceByIndex(cm.IfIndex)
					if err != nil {
						continue
					}
				}

				if n > 0 {
					m := new(dns.Msg)
					if err := m.Unpack(buf); err == nil && !shouldIgnore(m) {
						ch <- &Request{m, udpAddr, iface}
					}
				}
			}
		}()
	}
}

func (c *mdnsConn) sendQuery(m *dns.Msg, iface *net.Interface) error {
	sanitizeQuery(m)

	return c.writeMsg(m, iface)
}

func (c *mdnsConn) sendResponse(m *dns.Msg, iface *net.Interface) error {
	sanitizeResponse(m)

	return c.writeMsg(m, iface)
}

func (c *mdnsConn) sendResponseTo(m *dns.Msg, iface *net.Interface, addr *net.UDPAddr) error {
	sanitizeResponse(m)

	return c.writeMsgTo(m, iface, addr)
}

func (c *mdnsConn) writeMsg(m *dns.Msg, iface *net.Interface) error {
	var err error
	if c.ipv4 != nil {
		err = c.writeMsgTo(m, iface, AddrIPv4LinkLocalMulticast)
	}

	if c.ipv6 != nil {
		err = c.writeMsgTo(m, iface, AddrIPv6LinkLocalMulticast)
	}

	return err
}

func (c *mdnsConn) writeMsgTo(m *dns.Msg, iface *net.Interface, addr *net.UDPAddr) error {
	sanitizeMsg(m)

	var err error
	if c.ipv4 != nil && addr.IP.To4() != nil {
		if out, err := m.Pack(); err == nil {
			var ctrl *ipv4.ControlMessage
			if iface != nil {
				ctrl = &ipv4.ControlMessage{
					IfIndex: iface.Index,
				}
			}
			_, err = c.ipv4.WriteTo(out, ctrl, addr)
		}
	}

	if c.ipv6 != nil && addr.IP.To4() == nil {
		if out, err := m.Pack(); err == nil {
			var ctrl *ipv6.ControlMessage
			if iface != nil {
				ctrl = &ipv6.ControlMessage{
					IfIndex: iface.Index,
				}
			}
			_, err = c.ipv6.WriteTo(out, ctrl, addr)
		}
	}

	return err
}

func shouldIgnore(m *dns.Msg) bool {
	if m.Opcode != 0 {
		return true
	}

	if m.Rcode != 0 {
		return true
	}

	return false
}

func sanitizeResponse(m *dns.Msg) {
	if m.Question != nil && len(m.Question) > 0 {
		log.Info.Println("Multicast DNS responses MUST NOT contain any questions in the Question Section.  (RFC6762 6)")
		m.Question = nil
	}

	if !m.Response {
		log.Info.Println("In response messages the QR bit MUST be one (RFC6762 18.2)")
		m.Response = true
	}

	if !m.Authoritative {
		log.Info.Println("AA Bit bit MUST be set to one in response messages (RFC6762 18.4)")
		m.Authoritative = true
	}

	if m.Truncated {
		log.Info.Println("In multicast response messages, the TC bit MUST be zero on transmission. (RFC6762 18.5)")
		m.Truncated = false
	}
}

func sanitizeQuery(m *dns.Msg) {
	if m.Response {
		log.Info.Println("In query messages the QR bit MUST be zero (RFC6762 18.2)")
		m.Response = false
	}

	if m.Authoritative {
		log.Info.Println("AA Bit MUST be zero in query messages (RFC6762 18.4)")
		m.Authoritative = false
	}
}

func sanitizeMsg(m *dns.Msg) {
	if m.Opcode != 0 {
		log.Info.Println("In both multicast query and multicast response messages, the OPCODE MUST be zero on transmission (RFC6762 18.3)")
		m.Opcode = 0
	}

	if m.RecursionDesired {
		log.Info.Println("In both multicast query and multicast response messages, the Recursion Available bit MUST be zero on transmission. (RFC6762 18.7)")
		m.RecursionDesired = false
	}

	if m.Zero {
		log.Info.Println("In both query and response messages, the Zero bit MUST be zero on transmission (RFC6762 18.8)")
		m.Zero = false
	}

	if m.AuthenticatedData {
		log.Info.Println("In both multicast query and multicast response messages, the Authentic Data bit MUST be zero on transmission (RFC6762 18.9)")
		m.AuthenticatedData = false
	}

	if m.CheckingDisabled {
		log.Info.Println("In both multicast query and multicast response messages, the Checking Disabled bit MUST be zero on transmission (RFC6762 18.10)")
		m.CheckingDisabled = false
	}

	if m.Rcode != 0 {
		log.Info.Println("In both multicast query and multicast response messages, the Response Code MUST be zero on transmission. (RFC6762 18.11)")
		m.Rcode = 0
	}
}

func first(errs ...error) error {
	for _, err := range errs {
		if err != nil {
			return err
		}
	}

	return nil
}

// Sets the Top Bit of rrclass for all answer records (except PTR) to trigger a cache flush in the receivers.
func setAnswerCacheFlushBit(msg *dns.Msg) {
	// From RFC6762
	//    The most significant bit of the rrclass for a record in the Answer
	//    Section of a response message is the Multicast DNS cache-flush bit
	//    and is discussed in more detail below in Section 10.2, "Announcements
	//    to Flush Outdated Cache Entries".
	for _, a := range msg.Answer {
		switch a.(type) {
		case *dns.PTR:
			continue
		default:
			a.Header().Class |= (1 << 15)
		}
	}
}

// Sets the Top Bit of class to indicate the unicast responses are preferred for this question.
func setQuestionUnicast(q *dns.Question) {
	q.Qclass |= (1 << 15)
}

// Returns true if q requires unicast responses.
func isUnicastQuestion(q dns.Question) bool {
	// From RFC6762
	// 18.12.  Repurposing of Top Bit of qclass in Question Section
	//
	//    In the Question Section of a Multicast DNS query, the top bit of the
	//    qclass field is used to indicate that unicast responses are preferred
	//    for this particular question.  (See Section 5.4.)
	return q.Qclass&(1<<15) != 0
}
<|endoftext|>


<commit_before>package rpg

import (
	"fmt"
	"image/color"
	"log"
	"os"
	"time"

	"golang.org/x/image/colornames"

	"github.com/faiface/pixel"
	"github.com/faiface/pixel/imdraw"
	"github.com/faiface/pixel/pixelgl"
)

type Button struct {
	Name  string
	Frame pixel.Rect
}

var version = "0.0.92"

func TitleMenu(w *World, win *pixelgl.Window) {
	title := NewText(40)
	text := NewText(24)
	dot := pixel.V(30, 400)
	title.Dot = dot
	title.Orig = title.Dot
	dot = pixel.V(-100, -300)
	text.Dot = dot
	text.Orig = text.Dot

	fmt.Fprintln(text, "https://github.com/aerth/rpg")
	fmt.Fprintf(title, "AERPG v%s\nPRESS ENTER", version)

	win.Clear(colornames.Black)

	win.SetTitle("AERPG (https://github.com/aerth/rpg)")
	log.Println("AERPG (https://github.com/aerth/rpg)")

	tick := time.Tick(time.Second)
	var frames = 0
	for !win.Closed() {

		win.Clear(colornames.Black)
		title.Color = RandomColor()
		text.Draw(win, pixel.IM)
		title.Draw(win, pixel.IM)
		frames++
		if win.JustPressed(pixelgl.KeyEscape) || win.JustPressed(pixelgl.KeyQ) {
			w.Exit(0)
		}
		if win.JustPressed(pixelgl.MouseButtonLeft) || win.JustPressed(pixelgl.KeyEnter) {
			break
		}
		win.Update()
		select {
		case <-tick:
			win.SetTitle(fmt.Sprintf("AERPG (https://github.com/aerth/rpg) %v fps", frames))
			frames = 0
		default:

		}
	}
	log.Println("thanks for playing")
}

var textmatrix = pixel.IM.Moved(pixel.V(10, 580))

func (w *World) IsButton(buttons []Button, point pixel.Vec) (Button, func(win pixel.Target, world *World), bool) {

	for _, button := range buttons {
		if button.Frame.Contains(point) {
			switch button.Name {
			case "manastorm":
				return button, func(win pixel.Target, world *World) {
					world.Action(w.Char, w.Char.Rect.Center(), ManaStorm)
				}, true
			default:
				return button, func(win pixel.Target, world *World) {
					world.Message(fmt.Sprintf("Bad button %s", point))
				}, true
			case "reset":
				return button, func(win pixel.Target, world *World) {
					world.Reset()
					//					world.Char.Inventory = []Item{}
				}, true

			}
		}
	}

	return Button{}, nil, false
}

func (w *World) Exit(code int) {
	os.Exit(code)
}

func DrawBar(imd *imdraw.IMDraw, color color.RGBA, cur, max float64, rect pixel.Rect) {
	imd.Color = color
	percent := cur / max
	one := rect.Min
	one.Y++
	imd.Push(rect.Min, rect.Max)
	imd.Rectangle(1)
	pt := pixel.V(rect.Max.X*percent, rect.Max.Y)
	imd.Push(rect.Min, pt)
	imd.Rectangle(0)

}

func (c *Character) DrawBars(target pixel.Target, bounds pixel.Rect) {
	var barheight = 10.00
	var startY = 50.00
	imd := imdraw.New(nil)
	xp := float64(c.Stats.XP)
	next := float64(c.NextLevel())
	rect := bounds
	rect.Min.Y = startY
	rect.Max.Y = rect.Min.Y + barheight
	DrawBar(imd, colornames.Red, float64(c.Health), float64(255), rect)
	DrawBar(imd, colornames.Blue, float64(c.Mana), 255.00, rect.Moved(pixel.V(0, barheight+1)))
	DrawBar(imd, colornames.Purple, xp, next, rect.Moved(pixel.V(0, (barheight*2)+1)))
	imd.Draw(target)
}
<commit_msg>add Version func<commit_after>package rpg

import (
	"fmt"
	"image/color"
	"log"
	"os"
	"time"

	"golang.org/x/image/colornames"

	"github.com/faiface/pixel"
	"github.com/faiface/pixel/imdraw"
	"github.com/faiface/pixel/pixelgl"
)

type Button struct {
	Name  string
	Frame pixel.Rect
}

var version = "0.0.93"

func Version() string {
	return "AERPG " + version
}

func TitleMenu(w *World, win *pixelgl.Window) {
	title := NewText(40)
	text := NewText(24)
	dot := pixel.V(30, 400)
	title.Dot = dot
	title.Orig = title.Dot
	dot = pixel.V(-100, -300)
	text.Dot = dot
	text.Orig = text.Dot

	fmt.Fprintln(text, "https://github.com/aerth/rpg")
	fmt.Fprintf(title, "AERPG v%s\nPRESS ENTER", version)

	win.Clear(colornames.Black)

	win.SetTitle("AERPG (https://github.com/aerth/rpg)")
	log.Println("AERPG (https://github.com/aerth/rpg)")

	tick := time.Tick(time.Second)
	var frames = 0
	for !win.Closed() {

		win.Clear(colornames.Black)
		title.Color = RandomColor()
		text.Draw(win, pixel.IM)
		title.Draw(win, pixel.IM)
		frames++
		if win.JustPressed(pixelgl.KeyEscape) || win.JustPressed(pixelgl.KeyQ) {
			w.Exit(0)
		}
		if win.JustPressed(pixelgl.MouseButtonLeft) || win.JustPressed(pixelgl.KeyEnter) {
			break
		}
		win.Update()
		select {
		case <-tick:
			win.SetTitle(fmt.Sprintf("AERPG (https://github.com/aerth/rpg) %v fps", frames))
			frames = 0
		default:

		}
	}
	log.Println("thanks for playing")
}

var textmatrix = pixel.IM.Moved(pixel.V(10, 580))

func (w *World) IsButton(buttons []Button, point pixel.Vec) (Button, func(win pixel.Target, world *World), bool) {

	for _, button := range buttons {
		if button.Frame.Contains(point) {
			switch button.Name {
			case "manastorm":
				return button, func(win pixel.Target, world *World) {
					world.Action(w.Char, w.Char.Rect.Center(), ManaStorm)
				}, true
			default:
				return button, func(win pixel.Target, world *World) {
					world.Message(fmt.Sprintf("Bad button %s", point))
				}, true
			case "reset":
				return button, func(win pixel.Target, world *World) {
					world.Reset()
					//					world.Char.Inventory = []Item{}
				}, true

			}
		}
	}

	return Button{}, nil, false
}

func (w *World) Exit(code int) {
	os.Exit(code)
}

func DrawBar(imd *imdraw.IMDraw, color color.RGBA, cur, max float64, rect pixel.Rect) {
	imd.Color = color
	percent := cur / max
	one := rect.Min
	one.Y++
	imd.Push(rect.Min, rect.Max)
	imd.Rectangle(1)
	pt := pixel.V(rect.Max.X*percent, rect.Max.Y)
	imd.Push(rect.Min, pt)
	imd.Rectangle(0)

}

func (c *Character) DrawBars(target pixel.Target, bounds pixel.Rect) {
	var barheight = 10.00
	var startY = 50.00
	imd := imdraw.New(nil)
	xp := float64(c.Stats.XP)
	next := float64(c.NextLevel())
	rect := bounds
	rect.Min.Y = startY
	rect.Max.Y = rect.Min.Y + barheight
	DrawBar(imd, colornames.Red, float64(c.Health), float64(255), rect)
	DrawBar(imd, colornames.Blue, float64(c.Mana), 255.00, rect.Moved(pixel.V(0, barheight+1)))
	DrawBar(imd, colornames.Purple, xp, next, rect.Moved(pixel.V(0, (barheight*2)+1)))
	imd.Draw(target)
}
<|endoftext|>


<commit_before>package mitm

import (
	"crypto/tls"
	"errors"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"sync"
	"time"
)

// Proxy is a forward proxy that substitutes its own certificate
// for incoming TLS connections in place of the upstream server's
// certificate.
type Proxy struct {
	// Wrap must be a function that wraps the upstream handler however it
	// wants.
	Wrap func(upstream http.Handler) http.Handler

	DialTLS func(network, addr string, c *tls.Config) (*tls.Conn, error)

	CA *tls.Certificate

	TLSServerConfig *tls.Config

	TLSClientConfig *tls.Config

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Println("handle", r.URL)
	if r.Method == "CONNECT" {
		p.serveConnect(w, r)
		return
	}
	rp := &httputil.ReverseProxy{
		Director:      httpDirector,
		FlushInterval: p.FlushInterval,
	}
	p.Wrap(rp).ServeHTTP(w, r)
}

func (p *Proxy) serveConnect(w http.ResponseWriter, r *http.Request) {
	var (
		err   error
		sconn *tls.Conn
		a     = dnsName(r.Host)
	)

	if len(a) == 0 {
		log.Println("cannot determine cert name for " + r.Host)
		http.Error(w, "no upstream", 503)
		return
	}

	provisionalCert, err := p.cert(a)
	if err != nil {
		log.Println("cert", err)
		http.Error(w, "no upstream", 503)
		return
	}

	sConfig := new(tls.Config)
	if p.TLSServerConfig != nil {
		*sConfig = *p.TLSServerConfig
	}
	sConfig.Certificates = []tls.Certificate{*provisionalCert}
	sConfig.GetCertificate = func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		cConfig := new(tls.Config)
		if p.TLSClientConfig != nil {
			*cConfig = *p.TLSClientConfig
		}
		cConfig.ServerName = hello.ServerName
		sconn, err = tls.Dial("tcp", r.Host, cConfig)
		if err != nil {
			log.Println("dial", r.Host, err)
			return nil, err
		}
		return p.cert([]string{hello.ServerName})
	}

	cconn, err := handshake(w, sConfig)
	if err != nil {
		log.Println("handshake", r.Host, err)
		return
	}
	defer cconn.Close()
	if sconn == nil {
		log.Println("could not determine cert name for " + r.Host)
		return
	}
	defer sconn.Close()

	od := &oneShotDialer{c: sconn}
	rp := &httputil.ReverseProxy{
		Director:      httpsDirector,
		Transport:     &http.Transport{DialTLS: od.dial},
		FlushInterval: p.FlushInterval,
	}

	ch := make(chan int)
	wc := &onCloseConn{cconn, func() { ch <- 0 }}
	http.Serve(&oneShotListener{wc}, p.Wrap(rp))
	<-ch
}

func (p *Proxy) cert(names []string) (*tls.Certificate, error) {
	return genCert(p.CA, names)
}

var okHeader = []byte("HTTP/1.1 200 OK\r\n\r\n")

func handshake(w http.ResponseWriter, config *tls.Config) (net.Conn, error) {
	raw, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		http.Error(w, "no upstream", 503)
		return nil, err
	}
	if _, err = raw.Write(okHeader); err != nil {
		raw.Close()
		return nil, err
	}
	conn := tls.Server(raw, config)
	err = conn.Handshake()
	if err != nil {
		conn.Close()
		raw.Close()
		return nil, err
	}
	return conn, nil
}

func httpDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "http"
}

func httpsDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "https"
}

// dnsName returns the DNS name in addr, if any.
func dnsName(addr string) []string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return nil
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return nil
	}
	return []string{host}
}

// namesOnCert returns the dns names
// in the peer's presented cert.
func namesOnCert(conn *tls.Conn) []string {
	// TODO(kr): handle IP addr SANs.
	c := conn.ConnectionState().PeerCertificates[0]
	if len(c.DNSNames) > 0 {
		// If Subject Alt Name is given,
		// we ignore the common name.
		// This matches behavior of crypto/x509.
		return c.DNSNames
	}
	return []string{c.Subject.CommonName}
}

type oneShotDialer struct {
	c  net.Conn
	mu sync.Mutex
}

func (d *oneShotDialer) dial(network, addr string) (net.Conn, error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.c == nil {
		return nil, errors.New("closed")
	}
	c := d.c
	d.c = nil
	return c, nil
}

type oneShotListener struct {
	c net.Conn
}

func (l *oneShotListener) Accept() (net.Conn, error) {
	if l.c == nil {
		return nil, errors.New("closed")
	}
	c := l.c
	l.c = nil
	return c, nil
}

func (l *oneShotListener) Close() error {
	return nil
}

func (l *oneShotListener) Addr() net.Addr {
	return l.c.LocalAddr()
}

type onCloseConn struct {
	net.Conn
	f func()
}

func (c *onCloseConn) Close() error {
	if c.f != nil {
		c.f()
		c.f = nil
	}
	return c.Conn.Close()
}
<commit_msg>mitm: remove defunct ParseIP<commit_after>package mitm

import (
	"crypto/tls"
	"errors"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"sync"
	"time"
)

// Proxy is a forward proxy that substitutes its own certificate
// for incoming TLS connections in place of the upstream server's
// certificate.
type Proxy struct {
	// Wrap must be a function that wraps the upstream handler however it
	// wants.
	Wrap func(upstream http.Handler) http.Handler

	DialTLS func(network, addr string, c *tls.Config) (*tls.Conn, error)

	CA *tls.Certificate

	TLSServerConfig *tls.Config

	TLSClientConfig *tls.Config

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Println("handle", r.URL)
	if r.Method == "CONNECT" {
		p.serveConnect(w, r)
		return
	}
	rp := &httputil.ReverseProxy{
		Director:      httpDirector,
		FlushInterval: p.FlushInterval,
	}
	p.Wrap(rp).ServeHTTP(w, r)
}

func (p *Proxy) serveConnect(w http.ResponseWriter, r *http.Request) {
	var (
		err   error
		sconn *tls.Conn
		name  = dnsName(r.Host)
	)

	if name == "" {
		log.Println("cannot determine cert name for " + r.Host)
		http.Error(w, "no upstream", 503)
		return
	}

	provisionalCert, err := p.cert(name)
	if err != nil {
		log.Println("cert", err)
		http.Error(w, "no upstream", 503)
		return
	}

	sConfig := new(tls.Config)
	if p.TLSServerConfig != nil {
		*sConfig = *p.TLSServerConfig
	}
	sConfig.Certificates = []tls.Certificate{*provisionalCert}
	sConfig.GetCertificate = func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		cConfig := new(tls.Config)
		if p.TLSClientConfig != nil {
			*cConfig = *p.TLSClientConfig
		}
		cConfig.ServerName = hello.ServerName
		sconn, err = tls.Dial("tcp", r.Host, cConfig)
		if err != nil {
			log.Println("dial", r.Host, err)
			return nil, err
		}
		return p.cert(hello.ServerName)
	}

	cconn, err := handshake(w, sConfig)
	if err != nil {
		log.Println("handshake", r.Host, err)
		return
	}
	defer cconn.Close()
	if sconn == nil {
		log.Println("could not determine cert name for " + r.Host)
		return
	}
	defer sconn.Close()

	od := &oneShotDialer{c: sconn}
	rp := &httputil.ReverseProxy{
		Director:      httpsDirector,
		Transport:     &http.Transport{DialTLS: od.dial},
		FlushInterval: p.FlushInterval,
	}

	ch := make(chan int)
	wc := &onCloseConn{cconn, func() { ch <- 0 }}
	http.Serve(&oneShotListener{wc}, p.Wrap(rp))
	<-ch
}

func (p *Proxy) cert(names ...string) (*tls.Certificate, error) {
	return genCert(p.CA, names)
}

var okHeader = []byte("HTTP/1.1 200 OK\r\n\r\n")

func handshake(w http.ResponseWriter, config *tls.Config) (net.Conn, error) {
	raw, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		http.Error(w, "no upstream", 503)
		return nil, err
	}
	if _, err = raw.Write(okHeader); err != nil {
		raw.Close()
		return nil, err
	}
	conn := tls.Server(raw, config)
	err = conn.Handshake()
	if err != nil {
		conn.Close()
		raw.Close()
		return nil, err
	}
	return conn, nil
}

func httpDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "http"
}

func httpsDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "https"
}

// dnsName returns the DNS name in addr, if any.
func dnsName(addr string) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return ""
	}
	return host
}

// namesOnCert returns the dns names
// in the peer's presented cert.
func namesOnCert(conn *tls.Conn) []string {
	// TODO(kr): handle IP addr SANs.
	c := conn.ConnectionState().PeerCertificates[0]
	if len(c.DNSNames) > 0 {
		// If Subject Alt Name is given,
		// we ignore the common name.
		// This matches behavior of crypto/x509.
		return c.DNSNames
	}
	return []string{c.Subject.CommonName}
}

type oneShotDialer struct {
	c  net.Conn
	mu sync.Mutex
}

func (d *oneShotDialer) dial(network, addr string) (net.Conn, error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.c == nil {
		return nil, errors.New("closed")
	}
	c := d.c
	d.c = nil
	return c, nil
}

type oneShotListener struct {
	c net.Conn
}

func (l *oneShotListener) Accept() (net.Conn, error) {
	if l.c == nil {
		return nil, errors.New("closed")
	}
	c := l.c
	l.c = nil
	return c, nil
}

func (l *oneShotListener) Close() error {
	return nil
}

func (l *oneShotListener) Addr() net.Addr {
	return l.c.LocalAddr()
}

type onCloseConn struct {
	net.Conn
	f func()
}

func (c *onCloseConn) Close() error {
	if c.f != nil {
		c.f()
		c.f = nil
	}
	return c.Conn.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
)

const (
	BankUpper = iota
	BankLower
)

type Mmc1 struct {
	RomBanks  [][]Word
	VromBanks [][]Word

	PrgBankCount  int
	ChrRomCount   int
	BatteryBacked bool
	Data          []byte

	Buffer        int
	BufferCounter uint
	PrgSwapBank   int
	PrgBankSize   int
	ChrBankSize   int
	Mirroring     int
}

func (m *Mmc1) Write(v Word, a int) {
	// If reset bit is set
	if v&0x80 != 0 {
		//fmt.Println("Resetting MMC")
		m.BufferCounter = 0
		m.Buffer = 0x0

		m.PrgSwapBank = BankLower
		m.PrgBankSize = Size16k
	} else {
		// Buffer the write
		m.Buffer = (m.Buffer & (0xFF - (0x1 << m.BufferCounter))) | ((int(v) & 0x1) << m.BufferCounter)
		m.BufferCounter++

		// If the buffer is filled
		if m.BufferCounter == 0x5 {
			m.SetRegister(m.RegisterNumber(a), m.Buffer)

			// Reset buffer
			m.BufferCounter = 0
			m.Buffer = 0
		}
	}
}

func (m *Mmc1) SetRegister(reg int, v int) {
	// fmt.Printf("Writing register %d: 0x%X\n", reg, v)
	switch reg {
	// Control register
	case 0:
		// Set mirroring
		tmp := v & 0x3

		if m.Mirroring != tmp {
			fmt.Printf("Mapper: MMC1\n  Mirroring: ")
			m.Mirroring = tmp
			switch m.Mirroring {
			case 0x0:
				ppu.Nametables.SetMirroring(MirroringSingleUpper)
				fmt.Println("Single Upper")
			case 0x1:
				ppu.Nametables.SetMirroring(MirroringSingleLower)
				fmt.Println("Single Lower")
			case 0x2:
				ppu.Nametables.SetMirroring(MirroringVertical)
				fmt.Println("Vertical")
			case 0x3:
				ppu.Nametables.SetMirroring(MirroringHorizontal)
				fmt.Println("Horizontal")
			}
		}

		switch (v >> 0x2) & 0x3 {
		case 0x0:
			fallthrough
		case 0x1:
			m.PrgBankSize = Size32k
			m.PrgSwapBank = BankLower
			//fmt.Printf("  PRG Swap Bank: Lower\n")
			//fmt.Printf("  PRG Bank Size: 32k\n")
		case 0x2:
			m.PrgBankSize = Size16k
			m.PrgSwapBank = BankUpper
			//fmt.Printf("  PRG Swap Bank: Upper\n")
			//fmt.Printf("  PRG Bank Size: 16k\n")
		case 0x3:
			m.PrgBankSize = Size16k
			m.PrgSwapBank = BankLower
			//fmt.Printf("  PRG Swap Bank: Lower\n")
			//fmt.Printf("  PRG Bank Size: 16k\n")
		}

		//fmt.Printf("  CHR Bank Size: ")
		// Set CHR bank size
		switch (v >> 0x4) & 0x1 {
		case 0x0:
			m.ChrBankSize = Size8k
			//fmt.Println("8k")
		case 0x1:
			m.ChrBankSize = Size4k
			//fmt.Println("4k")
		}
		// CHR Bank 0
	case 1:
		if m.ChrRomCount == 0 {
			return
		}

		//fmt.Printf("CHR Bank 0 -> ")
		// Select VROM at 0x0000
		switch m.ChrBankSize {
		case Size8k:
			// Swap 8k VROM (in 8k mode, ignore first bit D0)
			bank := v & 0xF

			//fmt.Printf("8k CHR write to: %d\n", bank)
			WriteVramBank(m.VromBanks, bank, 0x0000, Size4k)
			WriteVramBank(m.VromBanks, bank+1, 0x1000, Size4k)
		case Size4k:
			// Swap 4k VROM
			var bank int

			if v&0x10 == 0x10 {
				bank = (len(m.VromBanks) / 2) + (v & 0xF)
			} else {
				bank = v & 0xF
			}
			//fmt.Printf("4k CHR write to: %d\n", bank)
			WriteVramBank(m.VromBanks, bank, 0x0, Size4k)
		}
		// CHR Bank 1
	case 2:
		if m.ChrRomCount == 0 {
			return
		}

		//fmt.Printf("Value: 0x%X\n", v)
		//fmt.Printf("CHR Bank 1 -> ")
		// Select VROM bank at 0x1000, ignored in
		// 8k switching mode
		if m.ChrBankSize == Size4k {
			var bank int

			if v&0x10 == 0x10 {
				bank = (len(m.VromBanks) / 2) + (v & 0xF)
			} else {
				bank = v & 0xF
			}
			//fmt.Printf("4k CHR write to: %d\n", bank)
			WriteVramBank(m.VromBanks, bank, 0x1000, Size4k)
		}
		// PRG Bank
	case 3:
		if m.PrgBankCount >= 32 {
			//fmt.Println("32+ ROM banks")
		} else if m.PrgBankCount >= 16 {
			//fmt.Println("16+ ROM banks")
		}

		switch m.PrgBankSize {
		case Size32k:
			// Swap 32k ROM (in 32k mode, ignore first bit D0)
			bank := ((v >> 0x1) & 0x7) * 2
			fmt.Printf("32k write to: %d\n", bank/2)

			WriteRamBank(m.RomBanks, bank, 0x8000, Size16k)
			WriteRamBank(m.RomBanks, bank+1, 0xC000, Size16k)
		case Size16k:
			// Swap 16k ROM
			bank := v & 0xF

			if m.PrgSwapBank == BankUpper {
				// fmt.Printf("16k PRG UPPER write to: %d\n", bank)
				WriteRamBank(m.RomBanks, bank, 0xC000, Size16k)
			} else {
				// fmt.Printf("16k PRG LOWER write to: %d\n", bank)
				WriteRamBank(m.RomBanks, bank, 0x8000, Size16k)
			}
		}
	}
}

func (m *Mmc1) RegisterNumber(a int) int {
	switch {
	case a >= 0x8000 && a <= 0x9FFF:
		return 0
	case a >= 0xA000 && a <= 0xBFFF:
		return 1
	case a >= 0xC000 && a <= 0xDFFF:
		return 2
	}

	return 3
}
<commit_msg>Cleanup logging<commit_after>package main

import (
	"fmt"
)

const (
	BankUpper = iota
	BankLower
)

type Mmc1 struct {
	RomBanks  [][]Word
	VromBanks [][]Word

	PrgBankCount  int
	ChrRomCount   int
	BatteryBacked bool
	Data          []byte

	Buffer        int
	BufferCounter uint
	PrgSwapBank   int
	PrgBankSize   int
	ChrBankSize   int
	Mirroring     int
}

func (m *Mmc1) Write(v Word, a int) {
	// If reset bit is set
	if v&0x80 != 0 {
		m.BufferCounter = 0
		m.Buffer = 0x0

		m.PrgSwapBank = BankLower
		m.PrgBankSize = Size16k
	} else {
		// Buffer the write
		m.Buffer = (m.Buffer & (0xFF - (0x1 << m.BufferCounter))) | ((int(v) & 0x1) << m.BufferCounter)
		m.BufferCounter++

		// If the buffer is filled
		if m.BufferCounter == 0x5 {
			m.SetRegister(m.RegisterNumber(a), m.Buffer)

			// Reset buffer
			m.BufferCounter = 0
			m.Buffer = 0
		}
	}
}

func (m *Mmc1) SetRegister(reg int, v int) {
	switch reg {
	// Control register
	case 0:
		// Set mirroring
		tmp := v & 0x3

		if m.Mirroring != tmp {
			m.Mirroring = tmp
			switch m.Mirroring {
			case 0x0:
				ppu.Nametables.SetMirroring(MirroringSingleUpper)
			case 0x1:
				ppu.Nametables.SetMirroring(MirroringSingleLower)
			case 0x2:
				ppu.Nametables.SetMirroring(MirroringVertical)
			case 0x3:
				ppu.Nametables.SetMirroring(MirroringHorizontal)
			}
		}

		switch (v >> 0x2) & 0x3 {
		case 0x0:
			fallthrough
		case 0x1:
			m.PrgBankSize = Size32k
			m.PrgSwapBank = BankLower
		case 0x2:
			m.PrgBankSize = Size16k
			m.PrgSwapBank = BankUpper
		case 0x3:
			m.PrgBankSize = Size16k
			m.PrgSwapBank = BankLower
		}

		// Set CHR bank size
		switch (v >> 0x4) & 0x1 {
		case 0x0:
			m.ChrBankSize = Size8k
		case 0x1:
			m.ChrBankSize = Size4k
		}
		// CHR Bank 0
	case 1:
		if m.ChrRomCount == 0 {
			return
		}

		// Select VROM at 0x0000
		switch m.ChrBankSize {
		case Size8k:
			// Swap 8k VROM (in 8k mode, ignore first bit D0)
			bank := v & 0xF

			WriteVramBank(m.VromBanks, bank, 0x0000, Size4k)
			WriteVramBank(m.VromBanks, bank+1, 0x1000, Size4k)
		case Size4k:
			// Swap 4k VROM
			var bank int

			if v&0x10 == 0x10 {
				bank = (len(m.VromBanks) / 2) + (v & 0xF)
			} else {
				bank = v & 0xF
			}
			WriteVramBank(m.VromBanks, bank, 0x0, Size4k)
		}
		// CHR Bank 1
	case 2:
		if m.ChrRomCount == 0 {
			return
		}

		// Select VROM bank at 0x1000, ignored in
		// 8k switching mode
		if m.ChrBankSize == Size4k {
			var bank int

			if v&0x10 == 0x10 {
				bank = (len(m.VromBanks) / 2) + (v & 0xF)
			} else {
				bank = v & 0xF
			}
			WriteVramBank(m.VromBanks, bank, 0x1000, Size4k)
		}
		// PRG Bank
	case 3:
		if m.PrgBankCount >= 32 {
			//fmt.Println("32+ ROM banks")
		} else if m.PrgBankCount >= 16 {
			//fmt.Println("16+ ROM banks")
		}

		switch m.PrgBankSize {
		case Size32k:
			// Swap 32k ROM (in 32k mode, ignore first bit D0)
			bank := ((v >> 0x1) & 0x7) * 2
			fmt.Printf("32k write to: %d\n", bank/2)

			WriteRamBank(m.RomBanks, bank, 0x8000, Size16k)
			WriteRamBank(m.RomBanks, bank+1, 0xC000, Size16k)
		case Size16k:
			// Swap 16k ROM
			bank := v & 0xF

			if m.PrgSwapBank == BankUpper {
				WriteRamBank(m.RomBanks, bank, 0xC000, Size16k)
			} else {
				WriteRamBank(m.RomBanks, bank, 0x8000, Size16k)
			}
		}
	}
}

func (m *Mmc1) RegisterNumber(a int) int {
	switch {
	case a >= 0x8000 && a <= 0x9FFF:
		return 0
	case a >= 0xA000 && a <= 0xBFFF:
		return 1
	case a >= 0xC000 && a <= 0xDFFF:
		return 2
	}

	return 3
}
<|endoftext|>


<commit_before>// myLG is command line looking glass that written with Go language
// it tries from its own icmp and external looking glasses tools
package main

import (
	"errors"
	"github.com/briandowns/spinner"
	"net"
	"regexp"
	"strings"
	"time"

	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/http/ping"
	"github.com/mehrdadrad/mylg/icmp"
	"github.com/mehrdadrad/mylg/lg"
	"github.com/mehrdadrad/mylg/ns"
	"github.com/mehrdadrad/mylg/peeringdb"
	"github.com/mehrdadrad/mylg/ripe"
	"github.com/mehrdadrad/mylg/scan"
)

// Provider represents looking glass
type Provider interface {
	Set(host, version string)
	GetDefaultNode() string
	GetNodes() []string
	ChangeNode(node string)
	Ping() (string, error)
	Trace() chan string
	BGP() chan string
}

// Whois represents whois providers
type Whois interface {
	Set(r string)
	GetData() bool
	PrettyPrint()
}

const (
	version = "0.1.8"
)

var (
	// register looking glass hosts
	providers = map[string]Provider{"telia": new(lg.Telia), "level3": new(lg.Level3), "cogent": new(lg.Cogent)}
	whois     = map[string]Whois{"asn": new(ripe.ASN), "prefix": new(ripe.Prefix)}
	pNames    = providerNames()
	rep       = make(chan string, 1)
	req       = make(chan string, 1)
	nxt       = make(chan struct{}, 1)
	spin      = spinner.New(spinner.CharSets[26], 220*time.Millisecond)
	args      string
	prompt    string
	cPName    string
	nsr       *ns.Request
	c         *cli.Readline
)

// providerName
func providerNames() []string {
	pNames := []string{}
	for p := range providers {
		pNames = append(pNames, p)
	}
	return pNames
}

// validateProvider
func validateProvider(p string) (string, error) {
	pNames := []string{}
	match, _ := regexp.MatchString("("+strings.Join(pNames, "|")+")", p)
	p = strings.ToLower(p)
	if match {
		return p, nil
	}
	return "", errors.New("provider not support")

}
func init() {
	// initialize cli
	c = cli.Init("local", version)
	go c.Run(req, nxt)
	// initialize name server
	nsr = ns.NewRequest()
	go nsr.Init()
	// set default provider
	cPName = "local"
}

func main() {
	var (
		request string
		loop    = true
	)

	r, _ := regexp.Compile(`(ping|trace|bgp|lg|ns|dig|whois|peering|scan|hping|connect|node|local|mode|help|exit|quit)\s{0,1}(.*)`)

	for loop {
		select {
		case request, loop = <-req:
			if !loop {
				break
			}
			if len(request) < 1 {
				c.Next()
				continue
			}
			subReq := r.FindStringSubmatch(request)
			if len(subReq) == 0 {
				println("syntax error")
				c.Next()
				continue
			}
			prompt = c.GetPrompt()
			args = strings.TrimSpace(subReq[2])
			cmd := strings.TrimSpace(subReq[1])
			switch {
			case cmd == "hping" && cPName == "local":
				hping()
			case cmd == "ping" && cPName == "local":
				pingLocal()
			case cmd == "ping":
				pingLG()
			case cmd == "trace":
				trace()
			case cmd == "bgp":
				BGP()
			case cmd == "dig":
				nsr.Dig(args)
			case cmd == "node":
				node()
			case cmd == "local":
				nsr.Local()
				cPName = "local"
				c.SetPrompt(cPName)
			case cmd == "connect":
				connect()
			case cmd == "lg":
				c.SetPrompt("lg")
				c.UpdateCompleter("connect", pNames)
			case cmd == "ns":
				c.UpdateCompleter("connect", nsr.CountryList())
				c.UpdateCompleter("node", []string{})
				c.SetPrompt("ns")
			case cmd == "whois":
				whoIs()
			case cmd == "peering":
				peeringdb.Search(args)
			case cmd == "scan":
				scanPorts()
			case cmd == "mode":
				mode()
			case cmd == "help":
				c.Help()
			case cmd == "exit", cmd == "quit":
				c.Close(nxt)
				close(req)
			}
			// next line
			c.Next()
		}
	}
}

// node handles node cmd
func node() {
	switch {
	case strings.HasPrefix(prompt, "lg"):
		if _, ok := providers[cPName]; ok {
			providers[cPName].ChangeNode(args)
			c.UpdatePromptN(args, 3)
		} else {
			println("the specified node doesn't support")
		}
	case strings.HasPrefix(prompt, "ns"):
		if !nsr.ChkNode(args) {
			println("error: argument is not valid")
		} else {
			c.UpdatePromptN(args, 3)
		}
	}
}

// connect handles connect cmd
func connect() {
	var (
		pName string
		err   error
	)
	switch {
	case strings.HasPrefix(prompt, "lg"):
		if pName, err = validateProvider(args); err != nil {
			println("provider not available")
			c.Next()
			return
		}
		cPName = pName
		if _, ok := providers[cPName]; ok {
			c.UpdatePromptN(cPName+"/"+providers[cPName].GetDefaultNode(), 2)
			go func() {
				c.UpdateCompleter("node", providers[cPName].GetNodes())
			}()
		} else {
			println("it doesn't support")
		}
	case strings.HasPrefix(prompt, "ns"):
		if !nsr.ChkCountry(args) {
			println("error: argument is not valid")
		} else {
			c.SetPrompt("ns/" + args)
			c.UpdateCompleter("node", nsr.NodeList())
		}
	}
}

// whoIs tries to get whois information
// ASN and prefix/ip information
func whoIs() {
	if ripe.IsASN(args) {
		whois["asn"].Set(args)
		whois["asn"].GetData()
		whois["asn"].PrettyPrint()
	} else {
		whois["prefix"].Set(args)
		whois["prefix"].GetData()
		whois["prefix"].PrettyPrint()
	}
}

// mode set editor mode
func mode() {
	if args == "vim" {
		c.SetVim()
	} else if args == "emacs" {
		c.SetEmacs()
	} else {
		println("the request mode doesn't support")
	}
}

// trace tries to trace from local and lg
func trace() {
	switch {
	case strings.HasPrefix(prompt, "local"):
		trace := icmp.Trace{}
		trace.Run(args)
	case strings.HasPrefix(prompt, "lg"):
		providers[cPName].Set(args, "ipv4")
		for l := range providers[cPName].Trace() {
			println(l)
		}
	}
}

// hping tries to ping a web server by http
func hping() {
	p, err := ping.NewPing(args, 5)
	if err != nil {
		println(err.Error())
	} else {
		p.Run()
	}
}

// pingLG tries to ping through a looking glass
func pingLG() {
	spin.Prefix = "please wait "
	spin.Start()
	providers[cPName].Set(args, "ipv4")
	m, err := providers[cPName].Ping()
	spin.Stop()
	if err != nil {
		println(err.Error())
	} else {
		println(m)
	}
}

// pingLocal tries to ping from local source ip
func pingLocal() {
	p := icmp.NewPing()
	ra, err := net.ResolveIPAddr("ip", args)
	if err != nil {
		println("cannot resolve", args, ": Unknown host")
		return
	}
	p.IP(ra.String())
	for n := 0; n < 4; n++ {
		p.Ping(rep)
		println(<-rep)
	}
}

// scanPorts tries to scan tcp/ip ports
func scanPorts() {
	scan, err := scan.NewScan(args)
	if err != nil {
		println(err.Error())
	} else {
		scan.Run()
	}
}

// BGP tries to get BGP lookup from a LG
func BGP() {
	if cPName == "local" {
		println("no provider selected")
		return
	}
	providers[cPName].Set(args, "ipv4")
	for l := range providers[cPName].BGP() {
		println(l)
	}
}
<commit_msg>optimized based on the recent changes<commit_after>// myLG is command line looking glass that written with Go language
// it tries from its own icmp and external looking glasses tools
package main

import (
	"errors"
	"github.com/briandowns/spinner"
	"regexp"
	"strings"
	"time"

	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/http/ping"
	"github.com/mehrdadrad/mylg/icmp"
	"github.com/mehrdadrad/mylg/lg"
	"github.com/mehrdadrad/mylg/ns"
	"github.com/mehrdadrad/mylg/peeringdb"
	"github.com/mehrdadrad/mylg/scan"
	"github.com/mehrdadrad/mylg/whois"
)

// Provider represents looking glass
type Provider interface {
	Set(host, version string)
	GetDefaultNode() string
	GetNodes() []string
	ChangeNode(node string)
	Ping() (string, error)
	Trace() chan string
	BGP() chan string
}

const (
	version = "0.1.8"
)

var (
	// register looking glass hosts
	providers = map[string]Provider{"telia": new(lg.Telia), "level3": new(lg.Level3), "cogent": new(lg.Cogent)}
	pNames    = providerNames()
	req       = make(chan string, 1)
	nxt       = make(chan struct{}, 1)
	spin      = spinner.New(spinner.CharSets[26], 220*time.Millisecond)
	args      string
	prompt    string
	cPName    string
	nsr       *ns.Request
	c         *cli.Readline
)

// providerName
func providerNames() []string {
	pNames := []string{}
	for p := range providers {
		pNames = append(pNames, p)
	}
	return pNames
}

// validateProvider
func validateProvider(p string) (string, error) {
	pNames := []string{}
	match, _ := regexp.MatchString("("+strings.Join(pNames, "|")+")", p)
	p = strings.ToLower(p)
	if match {
		return p, nil
	}
	return "", errors.New("provider not support")

}
func init() {
	// initialize cli
	c = cli.Init("local", version)
	go c.Run(req, nxt)
	// initialize name server
	nsr = ns.NewRequest()
	go nsr.Init()
	// set default provider
	cPName = "local"
}

func main() {
	var (
		request string
		loop    = true
	)

	r, _ := regexp.Compile(`(ping|trace|bgp|lg|ns|dig|whois|peering|scan|hping|connect|node|local|mode|help|exit|quit)\s{0,1}(.*)`)

	for loop {
		select {
		case request, loop = <-req:
			if !loop {
				break
			}
			if len(request) < 1 {
				c.Next()
				continue
			}
			subReq := r.FindStringSubmatch(request)
			if len(subReq) == 0 {
				println("syntax error")
				c.Next()
				continue
			}
			prompt = c.GetPrompt()
			args = strings.TrimSpace(subReq[2])
			cmd := strings.TrimSpace(subReq[1])
			switch {
			case cmd == "hping" && cPName == "local":
				hping()
			case cmd == "ping" && cPName == "local":
				pingLocal()
			case cmd == "ping":
				pingLG()
			case cmd == "trace":
				trace()
			case cmd == "bgp":
				BGP()
			case cmd == "dig":
				nsr.Dig(args)
			case cmd == "node":
				node()
			case cmd == "local":
				nsr.Local()
				cPName = "local"
				c.SetPrompt(cPName)
			case cmd == "connect":
				connect()
			case cmd == "lg":
				c.SetPrompt("lg")
				c.UpdateCompleter("connect", pNames)
			case cmd == "ns":
				c.UpdateCompleter("connect", nsr.CountryList())
				c.UpdateCompleter("node", []string{})
				c.SetPrompt("ns")
			case cmd == "whois":
				whois.Lookup(args)
			case cmd == "peering":
				peeringdb.Search(args)
			case cmd == "scan":
				scanPorts()
			case cmd == "mode":
				mode()
			case cmd == "help":
				c.Help()
			case cmd == "exit", cmd == "quit":
				c.Close(nxt)
				close(req)
			}
			// next line
			c.Next()
		}
	}
}

// node handles node cmd
func node() {
	switch {
	case strings.HasPrefix(prompt, "lg"):
		if _, ok := providers[cPName]; ok {
			providers[cPName].ChangeNode(args)
			c.UpdatePromptN(args, 3)
		} else {
			println("the specified node doesn't support")
		}
	case strings.HasPrefix(prompt, "ns"):
		if !nsr.ChkNode(args) {
			println("error: argument is not valid")
		} else {
			c.UpdatePromptN(args, 3)
		}
	}
}

// connect handles connect cmd
func connect() {
	var (
		pName string
		err   error
	)
	switch {
	case strings.HasPrefix(prompt, "lg"):
		if pName, err = validateProvider(args); err != nil {
			println("provider not available")
			c.Next()
			return
		}
		cPName = pName
		if _, ok := providers[cPName]; ok {
			c.UpdatePromptN(cPName+"/"+providers[cPName].GetDefaultNode(), 2)
			go func() {
				c.UpdateCompleter("node", providers[cPName].GetNodes())
			}()
		} else {
			println("it doesn't support")
		}
	case strings.HasPrefix(prompt, "ns"):
		if !nsr.ChkCountry(args) {
			println("error: argument is not valid")
		} else {
			c.SetPrompt("ns/" + args)
			c.UpdateCompleter("node", nsr.NodeList())
		}
	}
}

// mode set editor mode
func mode() {
	if args == "vim" {
		c.SetVim()
	} else if args == "emacs" {
		c.SetEmacs()
	} else {
		println("the request mode doesn't support")
	}
}

// trace tries to trace from local and lg
func trace() {
	switch {
	case strings.HasPrefix(prompt, "local"):
		trace := icmp.Trace{}
		trace.Run(args)
	case strings.HasPrefix(prompt, "lg"):
		providers[cPName].Set(args, "ipv4")
		for l := range providers[cPName].Trace() {
			println(l)
		}
	}
}

// hping tries to ping a web server by http
func hping() {
	p, err := ping.NewPing(args, 5)
	if err != nil {
		println(err.Error())
	} else {
		p.Run()
	}
}

// pingLG tries to ping through a looking glass
func pingLG() {
	spin.Prefix = "please wait "
	spin.Start()
	providers[cPName].Set(args, "ipv4")
	m, err := providers[cPName].Ping()
	spin.Stop()
	if err != nil {
		println(err.Error())
	} else {
		println(m)
	}
}

// pingLocal tries to ping from local source ip
func pingLocal() {
	p, err := icmp.NewPing(args)
	if err != nil {
		return
	}
	p.Run()
}

// scanPorts tries to scan tcp/ip ports
func scanPorts() {
	scan, err := scan.NewScan(args)
	if err != nil {
		println(err.Error())
	} else {
		scan.Run()
	}
}

// BGP tries to get BGP lookup from a LG
func BGP() {
	if cPName == "local" {
		println("no provider selected")
		return
	}
	providers[cPName].Set(args, "ipv4")
	for l := range providers[cPName].BGP() {
		println(l)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"github.com/briandowns/spinner"
	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/icmp"
	"github.com/mehrdadrad/mylg/lg"
	"net"
	"regexp"
	"strings"
	"time"
)

type Provider interface {
	Set(host, version string)
	GetDefaultNode() string
	GetNodes() map[string]string
	ChangeNode(node string)
	Ping() (string, error)
}

var providers = map[string]Provider{"telia": new(lg.Telia), "level3": new(lg.Level3)}

func validateProvider(p string) (string, error) {
	match, _ := regexp.MatchString("(telia|level3)", p)
	p = strings.ToLower(p)
	if match {
		return p, nil
	} else {
		return "", errors.New("provider not support")
	}
}

func main() {
	var (
		err    error
		cPName string = "local"
	)

	rep := make(chan string, 1)
	cmd := make(chan string, 1)
	nxt := make(chan struct{}, 1)

	c := cli.Init("local")
	go c.Run(cmd, nxt)

	r, _ := regexp.Compile(`(ping|connect|node|local) (.*)`)
	s := spinner.New(spinner.CharSets[0], 220*time.Millisecond)

	for {
		select {
		case req := <-cmd:
			subReq := r.FindStringSubmatch(req)
			if len(subReq) == 0 {
				println("syntax error")
				nxt <- struct{}{}
				continue
			}
			switch {
			case subReq[1] == "ping" && cPName == "local":
				p := icmp.NewPing()
				ra, err := net.ResolveIPAddr("ip", subReq[2])
				if err != nil {
					println("cannot resolve", subReq[2], ": Unknown host")
					nxt <- struct{}{}
					continue
				}
				p.IP(ra.String())
				for n := 0; n < 4; n++ {
					p.Ping(rep)
					println(<-rep)
				}
				nxt <- struct{}{}
			case subReq[1] == "ping":
				s.Prefix = "please wait "
				s.Start()
				providers[cPName].Set(subReq[2], "ipv4")
				m, _ := providers[cPName].Ping()
				s.Stop()
				println(m)
				nxt <- struct{}{}
			case subReq[1] == "node":
				providers[cPName].ChangeNode(subReq[2])
				c.SetPrompt(cPName + "/" + subReq[2])
				nxt <- struct{}{}
			case subReq[1] == "local":
				cPName = "local"
				c.SetPrompt(cPName)
				nxt <- struct{}{}
			case subReq[1] == "connect":
				var pName string
				if pName, err = validateProvider(subReq[2]); err != nil {
					println("provider not available")
					nxt <- struct{}{}
					continue
				}
				cPName = pName
				c.SetPrompt(cPName + "/" + providers[cPName].GetDefaultNode())
				nxt <- struct{}{}
			}
		}
	}
}
<commit_msg>changed spinning<commit_after>package main

import (
	"errors"
	"github.com/briandowns/spinner"
	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/icmp"
	"github.com/mehrdadrad/mylg/lg"
	"net"
	"regexp"
	"strings"
	"time"
)

type Provider interface {
	Set(host, version string)
	GetDefaultNode() string
	GetNodes() map[string]string
	ChangeNode(node string)
	Ping() (string, error)
}

var providers = map[string]Provider{"telia": new(lg.Telia), "level3": new(lg.Level3)}

func validateProvider(p string) (string, error) {
	match, _ := regexp.MatchString("(telia|level3)", p)
	p = strings.ToLower(p)
	if match {
		return p, nil
	} else {
		return "", errors.New("provider not support")
	}
}

func main() {
	var (
		err    error
		cPName string = "local"
	)

	rep := make(chan string, 1)
	cmd := make(chan string, 1)
	nxt := make(chan struct{}, 1)

	c := cli.Init("local")
	go c.Run(cmd, nxt)

	r, _ := regexp.Compile(`(ping|connect|node|local) (.*)`)
	s := spinner.New(spinner.CharSets[26], 220*time.Millisecond)

	for {
		select {
		case req := <-cmd:
			subReq := r.FindStringSubmatch(req)
			if len(subReq) == 0 {
				println("syntax error")
				nxt <- struct{}{}
				continue
			}
			switch {
			case subReq[1] == "ping" && cPName == "local":
				p := icmp.NewPing()
				ra, err := net.ResolveIPAddr("ip", subReq[2])
				if err != nil {
					println("cannot resolve", subReq[2], ": Unknown host")
					nxt <- struct{}{}
					continue
				}
				p.IP(ra.String())
				for n := 0; n < 4; n++ {
					p.Ping(rep)
					println(<-rep)
				}
				nxt <- struct{}{}
			case subReq[1] == "ping":
				s.Prefix = "please wait "
				s.Start()
				providers[cPName].Set(subReq[2], "ipv4")
				m, _ := providers[cPName].Ping()
				s.Stop()
				println(m)
				nxt <- struct{}{}
			case subReq[1] == "node":
				providers[cPName].ChangeNode(subReq[2])
				c.SetPrompt(cPName + "/" + subReq[2])
				nxt <- struct{}{}
			case subReq[1] == "local":
				cPName = "local"
				c.SetPrompt(cPName)
				nxt <- struct{}{}
			case subReq[1] == "connect":
				var pName string
				if pName, err = validateProvider(subReq[2]); err != nil {
					println("provider not available")
					nxt <- struct{}{}
					continue
				}
				cPName = pName
				c.SetPrompt(cPName + "/" + providers[cPName].GetDefaultNode())
				nxt <- struct{}{}
			}
		}
	}
}
<|endoftext|>


<commit_before>package cistern

import "errors"

const (
	IsRed       bool = true
	IsBlack     bool = false
	LeftRotate  bool = true
	RightRotate bool = false
)

// 定义树，默认只包含一个根节点
type Tree struct {
	root *TreeNode
}

func (t *Tree) insert(node *TreeNode) {

}

// 定义节点
type TreeNode struct {
	value               []byte
	key                 string
	color               bool
	left, right, parent *TreeNode
}

func (tn *TreeNode) getParent() *TreeNode {
	return tn.parent
}

// 树旋转，如果有根节点变动则返回变动后的根节点
func (tn *TreeNode) rotate(leftRotate bool) (*TreeNode, error) {
	var root *TreeNode
	if tn == nil {
		return root, nil
	}
	if leftRotate && tn.right == nil {
		return root, errors.New("right node must not be nil")
	}
	if !leftRotate && tn.left == nil {
		return root, errors.New("left node must not be nil")
	}
	parent := tn.parent
	var isLeft bool
	if parent != nil {
		isLeft = parent.left == tn
	}
	if leftRotate { // 左选
		grandson := tn.right.left
		tn.parent = tn.right
		tn.parent.left = tn
		tn.right = grandson
	} else { // 右旋
		grandson := tn.left.right
		tn.parent = tn.left
		tn.parent.right = tn
		tn.left = grandson
	}
	if parent == nil {
		tn.parent.parent = nil
		root = tn.parent
	} else {
		if isLeft {
			parent.left = tn.parent
		} else {
			parent.right = tn.parent
		}
		tn.parent.parent = parent

	}
	return root, nil
}
<commit_msg>在红黑树中插入节点<commit_after>package cistern

import "errors"

const (
	IsRed       bool = true
	IsBlack     bool = false
	LeftRotate  bool = true
	RightRotate bool = false
)

// 定义树，默认只包含一个根节点
type Tree struct {
	root *RBTreeNode
}

func (t *Tree) insert(current *RBTreeNode, key string, value []byte) {
	hash := hash(key)
	if hash < current.hash {
		if current.left == nil {
			node := NewNode(key, value, hash)
			current.left = node
			node.parent = current
			t.check(node)
		} else {
			t.insert(current.left, key, value)
		}
	} else if hash > current.hash {
		if current.right == nil {
			node := NewNode(key, value, hash)
			current.right = node
			node.parent = current
			t.check(node)
		} else {
			t.insert(current.right, key, value)
		}
	} else {
		current.value = value
	}
}

// 根据红黑树规则 校验和调整树形结构
func (t *Tree) check(node *RBTreeNode) {

}

// 定义节点
type RBTreeNode struct {
	hash                uint32
	key                 string
	value               []byte
	color               bool
	left, right, parent *RBTreeNode
}

func NewNode(key string, value []byte, hash uint32) (tNode *RBTreeNode) {
	tNode = new(RBTreeNode)
	tNode.key = key
	tNode.value = value
	tNode.hash = hash
	return
}

func (tn *RBTreeNode) getParent() *RBTreeNode {
	return tn.parent
}

// 树旋转，如果有根节点变动则返回变动后的根节点
func (tn *RBTreeNode) rotate(leftRotate bool) (*RBTreeNode, error) {
	var root *RBTreeNode
	if tn == nil {
		return root, nil
	}
	if leftRotate && tn.right == nil {
		return root, errors.New("right node must not be nil")
	}
	if !leftRotate && tn.left == nil {
		return root, errors.New("left node must not be nil")
	}
	parent := tn.parent
	var isLeft bool
	if parent != nil {
		isLeft = parent.left == tn
	}
	if leftRotate { // 左选
		grandson := tn.right.left
		tn.parent = tn.right
		tn.parent.left = tn
		tn.right = grandson
	} else { // 右旋
		grandson := tn.left.right
		tn.parent = tn.left
		tn.parent.right = tn
		tn.left = grandson
	}
	if parent == nil {
		tn.parent.parent = nil
		root = tn.parent
	} else {
		if isLeft {
			parent.left = tn.parent
		} else {
			parent.right = tn.parent
		}
		tn.parent.parent = parent

	}
	return root, nil
}

func hash(s string) uint32 {
	return 0 // TODO
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2013-2015 Oryx(ossrs)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

// The os defines:
//      bsd: darwin dragonfly freebsd nacl netbsd openbsd solaris
//      unix: bsd linux
//      server: unix plan9
//      posix: bsd linux windows
// All os by go:
//      server windows
//      posix plan9

package main

import (
	"flag"
	"fmt"
	"github.com/ossrs/go-oryx/app"
	"github.com/ossrs/go-oryx/core"
	"os"
)

func serve(svr *app.Server) int {
	if err := svr.PrepareLogger(); err != nil {
		core.Error.Println("prepare logger failed, err is", err)
		return -1
	}

	oryxMain(svr)

	core.Trace.Println("Copyright (c) 2013-2015 Oryx(ossrs)")
	core.Trace.Println(fmt.Sprintf("go-oryx/%v is advanced SRS, focus on realtime live streaming.", core.Version()))

	if err := svr.Initialize(); err != nil {
		core.Error.Println("initialize server failed, err is", err)
		return -1
	}

	if err := svr.Run(); err != nil {
		core.Error.Println("run server failed, err is", err)
		return -1
	}

	return 0
}

func main() {
	// the startup argv:
	//          -c conf/oryx.json
	//          --c conf/oryx.json
	//          -c=conf/oryx.json
	//          --c=conf/oryx.json
	var confFile = flag.String("c", "conf/oryx.json", "the config file.")

	flag.Usage = func(){
		fmt.Println(fmt.Sprintf("Usage: %v [--c=string] [-h|--help]", os.Args[0]))
		fmt.Println(fmt.Sprintf("	c, the config file path"))
		fmt.Println(fmt.Sprintf("	help, show this help and exit"))
		fmt.Println(fmt.Sprintf("For example:"))
		fmt.Println(fmt.Sprintf("	%v --c=conf/oryx.json", os.Args[0]))
	}
	flag.Parse()

	ret := func() int {
		svr := app.NewServer()
		defer svr.Close()

		if err := svr.ParseConfig(*confFile); err != nil {
			core.Error.Println("parse config from", *confFile, "failed, err is", err)
			return -1
		}

		return run(svr)
	}()

	os.Exit(ret)
}
<commit_msg>refine code of flag parse.<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2013-2015 Oryx(ossrs)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

// The os defines:
//      bsd: darwin dragonfly freebsd nacl netbsd openbsd solaris
//      unix: bsd linux
//      server: unix plan9
//      posix: bsd linux windows
// All os by go:
//      server windows
//      posix plan9

package main

import (
	"flag"
	"fmt"
	"github.com/ossrs/go-oryx/app"
	"github.com/ossrs/go-oryx/core"
	"os"
)

func serve(svr *app.Server) int {
	if err := svr.PrepareLogger(); err != nil {
		core.Error.Println("prepare logger failed, err is", err)
		return -1
	}

	oryxMain(svr)

	core.Trace.Println("Copyright (c) 2013-2015 Oryx(ossrs)")
	core.Trace.Println(fmt.Sprintf("go-oryx/%v is advanced SRS, focus on realtime live streaming.", core.Version()))

	if err := svr.Initialize(); err != nil {
		core.Error.Println("initialize server failed, err is", err)
		return -1
	}

	if err := svr.Run(); err != nil {
		core.Error.Println("run server failed, err is", err)
		return -1
	}

	return 0
}

func main() {
	// the startup argv:
	//          -c conf/oryx.json
	//          --c conf/oryx.json
	//          -c=conf/oryx.json
	//          --c=conf/oryx.json
	//          --conf=conf/oryx.json
	var confFile string
	flag.StringVar(&confFile, "c", "", "the config file")
	flag.StringVar(&confFile, "conf", "", "the config file")

	flag.Usage = func(){
		fmt.Println(fmt.Sprintf("Usage: %v [-c|--conf <filename>] [-?|-h|--help]", os.Args[0]))
		fmt.Println(fmt.Sprintf("	-c, --conf filename     : the config file path"))
		fmt.Println(fmt.Sprintf("	-?, -h, --help          : show this help and exit"))
		fmt.Println(fmt.Sprintf("For example:"))
		fmt.Println(fmt.Sprintf("	%v -c conf/oryx.json", os.Args[0]))
	}
	flag.Parse()

	// show help without conf.
	if len(confFile) == 0 {
		flag.Usage()
		os.Exit(0)
	}

	ret := func() int {
		svr := app.NewServer()
		defer svr.Close()

		if err := svr.ParseConfig(confFile); err != nil {
			core.Error.Println("parse config from", confFile, "failed, err is", err)
			return -1
		}

		return run(svr)
	}()

	os.Exit(ret)
}
<|endoftext|>


<commit_before>// The graph package implements a breadth-first-search-based pathfinding
// algorithm which runs faster than Dijkstra's algorithm in certain cases.
// This algorithm assumes that all edges have integral length. (Note that
// this should never provide accuracy issues. For a graph in which distances
// are stored as uint64, it is possible to measure edge-lengths up to 145
// trillion miles while maintaining accuracy of one inch.)
package graph

type Node interface {
	// Returns the number of edges
	// pointing away from this node
	Edges() int

	// Returns the node at the
	// end of the edge at index i
	Next(i int) Node

	// If the length of the edge
	// at index i is greater than 1,
	// decrement it by one and return
	// false. Else, return true.
	Dec(i int) bool

	// Return whether or not this
	// node's previous node has
	// been decided yet.
	PrevCheck() bool

	// Set this node's previous
	// node to n.
	PrevSet(n Node)

	/*
		This Prev function is nowhere
		used in this code. However, most
		reasonable uses of this package
		require such functionality on the
		part of the user, and so this is
		included to save them the headache
		of type assertions.
	*/
	// Return this node's
	// previous node
	Prev() Node

	// Returns whether or not
	// this node is the
	// destination node
	IsEnd() bool

	/*
		This IsStart function is nowhere
		used in this code. However, most
		reasonable uses of this package
		require such functionality on the
		part of the user, and so this is
		included to save them the headache
		of type assertions.
	*/
	// Returns whether or not
	// this node is the
	// start node
	IsStart() bool
}

type llNode struct {
	Node
	next *llNode
}

type ll struct {
	head *llNode
	tail *llNode
}

func (l *ll) add(n Node) {
	newNode := new(llNode)
	newNode.Node = n

	if l.head == nil {
		l.head = newNode
		l.tail = newNode
	} else {
		l.tail.next = newNode
		l.tail = newNode
	}
}

// Takes start and end nodes, and returns
// path length; return -1 if no path was found.
// Note that begining and ending nodes
// being the same will cause an infinite loop.
// It is an invariant that if every node in
// the graph is of the same type when FindPath
// is called, no new types will be added.
// Most importantly, any type assertions which
// would have been guaranted to succeed before
// FindPath retain the guarantee after FindPath
// is called.
func FindPath(start Node) int {
	// Set start's previous pointer
	// because nil is a sentinal value.
	// This is never actually used, so
	// it doesn't cause recursion problems
	start.PrevSet(start)

	listA := new(ll)
	listB := new(ll)
	listA.add(start)

	// Each iteration of BFS increases
	// path length by 1
	i := 1
END:
	for ; ; i++ {
		if listA.head == nil {
			return -1
		}
		for cur := listA.head; cur != nil; cur = cur.next {
			// Iterate through this node's adjacent nodes
			edges := cur.Node.Edges()
			added := false
			for j := 0; j < edges; j++ {
				next := cur.Node.Next(j)
				if cur.Node.Dec(j) && !next.PrevCheck() {
					next.PrevSet(cur.Node)
					if next.IsEnd() {
						break END
					}
					listB.add(next)
				} else if !added {
					listB.add(cur.Node)
					added = true
				}
			}
		}
		listA = listB
		listB = new(ll)
	}

	// Undo the change from
	// the beginning
	start.PrevSet(nil)
	return i
}
<commit_msg>Updated package name<commit_after>// The bfspath package implements a breadth-first-search-based pathfinding
// algorithm which runs faster than Dijkstra's algorithm in certain cases.
// This algorithm assumes that all edges have integral length. (Note that
// this should never provide accuracy issues. For a graph in which distances
// are stored as uint64, it is possible to measure edge-lengths up to 145
// trillion miles while maintaining accuracy of one inch.)
package bfspath

type Node interface {
	// Returns the number of edges
	// pointing away from this node
	Edges() int

	// Returns the node at the
	// end of the edge at index i
	Next(i int) Node

	// If the length of the edge
	// at index i is greater than 1,
	// decrement it by one and return
	// false. Else, return true.
	Dec(i int) bool

	// Return whether or not this
	// node's previous node has
	// been decided yet.
	PrevCheck() bool

	// Set this node's previous
	// node to n.
	PrevSet(n Node)

	/*
		This Prev function is nowhere
		used in this code. However, most
		reasonable uses of this package
		require such functionality on the
		part of the user, and so this is
		included to save them the headache
		of type assertions.
	*/
	// Return this node's
	// previous node
	Prev() Node

	// Returns whether or not
	// this node is the
	// destination node
	IsEnd() bool

	/*
		This IsStart function is nowhere
		used in this code. However, most
		reasonable uses of this package
		require such functionality on the
		part of the user, and so this is
		included to save them the headache
		of type assertions.
	*/
	// Returns whether or not
	// this node is the
	// start node
	IsStart() bool
}

type llNode struct {
	Node
	next *llNode
}

type ll struct {
	head *llNode
	tail *llNode
}

func (l *ll) add(n Node) {
	newNode := new(llNode)
	newNode.Node = n

	if l.head == nil {
		l.head = newNode
		l.tail = newNode
	} else {
		l.tail.next = newNode
		l.tail = newNode
	}
}

// Takes start and end nodes, and returns
// path length; return -1 if no path was found.
// Note that begining and ending nodes
// being the same will cause an infinite loop.
// It is an invariant that if every node in
// the graph is of the same type when FindPath
// is called, no new types will be added.
// Most importantly, any type assertions which
// would have been guaranted to succeed before
// FindPath retain the guarantee after FindPath
// is called.
func FindPath(start Node) int {
	// Set start's previous pointer
	// because nil is a sentinal value.
	// This is never actually used, so
	// it doesn't cause recursion problems
	start.PrevSet(start)

	listA := new(ll)
	listB := new(ll)
	listA.add(start)

	// Each iteration of BFS increases
	// path length by 1
	i := 1
END:
	for ; ; i++ {
		if listA.head == nil {
			return -1
		}
		for cur := listA.head; cur != nil; cur = cur.next {
			// Iterate through this node's adjacent nodes
			edges := cur.Node.Edges()
			added := false
			for j := 0; j < edges; j++ {
				next := cur.Node.Next(j)
				if cur.Node.Dec(j) && !next.PrevCheck() {
					next.PrevSet(cur.Node)
					if next.IsEnd() {
						break END
					}
					listB.add(next)
				} else if !added {
					listB.add(cur.Node)
					added = true
				}
			}
		}
		listA = listB
		listB = new(ll)
	}

	// Undo the change from
	// the beginning
	start.PrevSet(nil)
	return i
}
<|endoftext|>


<commit_before>package zmq4

/*
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <netdb.h>

char *my_get_peer_addr(int srcFd) {
    struct sockaddr_storage
        ss;
    socklen_t
        addrlen;
    int
        rc;
    char
        host [NI_MAXHOST];

    // get the remote endpoint
    addrlen = sizeof ss;
    rc = getpeername (srcFd, (struct sockaddr*) &ss, &addrlen);
    if (rc != 0)
        return NULL;

    rc = getnameinfo ((struct sockaddr*) &ss, addrlen, host, sizeof host, NULL, 0, NI_NUMERICHOST);
    if (rc != 0)
        return NULL;

    return strdup(host);
}
*/
import "C"

import (
	"unsafe"
)

// Given a socket fd, return the peer address.
func GetPeerAddr(fd int) (string, error) {
	s, err := C.my_get_peer_addr(C.int(fd))
	if s == nil {
		return "", errget(err)
	}
	defer C.free(unsafe.Pointer(s))
	return C.GoString(s), nil
}
<commit_msg>GetPeerAddr on Windows.<commit_after>package zmq4

/*
#include <stdlib.h>
#include <string.h>
#ifdef _WIN32
#include <Winsock2.h>
#include <Ws2tcpip.h>
#else
#include <sys/socket.h>
#include <netdb.h>
#endif

char *my_get_peer_addr(int srcFd) {
    struct sockaddr_storage
        ss;
    socklen_t
        addrlen;
    int
        rc;
    char
        host [NI_MAXHOST];

    // get the remote endpoint
    addrlen = sizeof ss;
    rc = getpeername (srcFd, (struct sockaddr*) &ss, &addrlen);
    if (rc != 0)
        return NULL;

    rc = getnameinfo ((struct sockaddr*) &ss, addrlen, host, sizeof host, NULL, 0, NI_NUMERICHOST);
    if (rc != 0)
        return NULL;

    return strdup(host);
}
*/
import "C"

import (
	"unsafe"
)

// Given a socket fd, return the peer address.
func GetPeerAddr(fd int) (string, error) {
	s, err := C.my_get_peer_addr(C.int(fd))
	if s == nil {
		return "", errget(err)
	}
	defer C.free(unsafe.Pointer(s))
	return C.GoString(s), nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/rcrowley/go-metrics"
)

var (
	output    = flag.String("o", "-", "output file")
	workers   = flag.Int("n", 4, "parallel connections")
	chunkSize = flag.Int("c", 100*1024, "chunk size")
	debug     = flag.Bool("d", false, "debug")
)

type chunk struct {
	offset int64
	size   int64
	file   string
}

type download struct {
	wg       sync.WaitGroup
	url      string
	size     int64
	startPos int64
	writer   io.Writer

	succeeded     counter
	downloaded    counter
	downloadMeter metrics.Meter
	progress      counter
}

type counter struct {
	v int64
}

func (c *counter) add(delta int64) {
	atomic.AddInt64(&c.v, delta)
}

func (c *counter) value() int64 {
	return atomic.LoadInt64(&c.v)
}

func newDownload(url string) (*download, error) {
	size, err := contentSize(url)
	if err != nil {
		return nil, err
	}

	return &download{
		url:           url,
		size:          size,
		writer:        os.Stdout,
		downloadMeter: metrics.NewMeter(),
	}, nil
}

// chunkFeeder generates chunk structures and sends them to the out channel.
// The chunk size is determined via commandline params. The last chunk is
// is sized accordingly.
func (d *download) chunkFeeder(out chan<- *chunk) {
	// when resuming a download. the chunk feeder won't emit chunks
	// that are already been written.
	d.succeeded.add(d.startPos)

	if *debug {
		log.Printf("Feeder is running")
	}
	for i := int64(0); ; i++ {
		size := int64(*chunkSize)
		off := i * size

		if off+size <= d.startPos {
			continue
		}
		if d.startPos-off > 0 && d.startPos-off < size {
			size -= d.startPos - off
			off = d.startPos
		}

		if off >= d.size {
			break
		}
		if off+size > d.size {
			size = off + size - d.size
		}
		if *debug {
			log.Printf("Enqueuing chunk %d size %d", off, size)
		}
		out <- &chunk{
			offset: off,
			size:   size}
	}
	if *debug {
		log.Printf("Feeder is closing")
	}
	close(out)
}

// chunkWorker task chunks from the in channel, fetches the chunk from
// the http resource into the buffer in the chunk structure.
// The completed chunk is then sent to the out channel.
// Any error causes to panic.
func (d *download) chunkWorker(n int, in <-chan *chunk, out chan<- *chunk) {
	if *debug {
		log.Println("Started chunk worker", n)
	}

	for ch := range in {
		if *debug {
			log.Printf("Chunk worker %d fetching chunk %d (%d)\n", n, ch.offset, ch.size)
		}

		basename := *output
		if *output == "-" {
			basename = "/tmp/pget.tmp"
		}

		offset := ch.offset
		size := ch.size

		var f *os.File
		var err error

		tmpName := fmt.Sprintf("%s.chunk-%d", basename, ch.offset)
		if st, err := os.Stat(tmpName); os.IsNotExist(err) {
			f, err = os.Create(tmpName)
			if err != nil {
				log.Fatalf("Cannot create temp file for chunk for offset %d, %s", ch.offset, err)
			}
		} else if err == nil {
			f, err = os.OpenFile(tmpName, os.O_APPEND, 0)
			if err != nil {
				log.Fatalf("Cannot open temp file for appending chunk for offset %d, %s", ch.offset, err)
			}
			offset += st.Size()
			size -= st.Size()
		} else {
			log.Fatal("Got some error", err)
		}

		var body io.ReadCloser
		if size < 1 {
			//fmt.Println("Fully downloaded chunk, skipping")
			body = ioutil.NopCloser(&bytes.Buffer{})
		} else {
			body, err = d.fetch(offset, size)
			if err != nil {
				log.Fatal(err)
			}
			// mark the whole chunk as succeeded
			// even if we continued an interrupted download.
			d.succeeded.add(ch.size)
		}

		// TODO(mkm) check errors
		io.Copy(f, body)
		body.Close()
		f.Close()

		ch.file = tmpName
		out <- ch
	}

	if *debug {
		log.Println("Finishing chunk worker", n)
	}
}

// chunkWriter reads completed chunks from the in channel,
// and writes them out in the correct order.
// Out of order chunks are kept in memory until the missing pieces
// are delivered.
func (d *download) chunkWriter(in <-chan *chunk) {
	readyChunks := make(map[int64]*chunk)
	outPos := d.startPos
	d.progress.add(d.startPos)

	for ch := range in {
		if *debug {
			log.Printf("Chunk writer processing completed chunk %d (%d)\n", ch.offset, ch.size)
		}
		readyChunks[ch.offset] = ch

		for {
			if c, ok := readyChunks[outPos]; ok {
				f, err := os.Open(c.file)
				if err != nil {
					log.Fatal(err)
				}
				n, err := io.Copy(d.writer, f)
				if err != nil {
					log.Fatal(err)
				}
				d.progress.add(n)

				f.Close()
				os.Remove(c.file)

				outPos += c.size
			} else {
				break
			}
		}
	}
}

// start
func start() error {
	url := flag.Arg(0)

	d, err := newDownload(url)
	if err != nil {
		return err
	}

	lastPos := int64(0)
	if *output != "-" {
		//f, err := os.Create(*output)
		f, err := os.OpenFile(*output, os.O_RDWR|os.O_CREATE, 0666)
		if err != nil {
			return err
		}
		defer f.Close()
		d.writer = f

		lastPos, err = f.Seek(0, 2)
		if err != nil {
			log.Fatal(err)
		}
	}
	d.startPos = lastPos

	if *debug {
		log.Println("Size", d.size)
		log.Println("Last pos", lastPos)
	}

	ich := make(chan *chunk, 2)
	och := make(chan *chunk, 2)

	// spawn the chunk feeder.
	go d.chunkFeeder(ich)

	// spawn the workers in such a way
	// that we know when they are all done.
	for i := 0; i < *workers; i++ {
		d.wg.Add(1)
		go func(i int) {
			d.chunkWorker(i, ich, och)
			d.wg.Done()
		}(i)
	}

	// when all the workers are gone
	// close the chunk writer input
	go func() {
		d.wg.Wait()
		close(och)
	}()

	done := make(chan struct{})
	go d.showProgress(done)

	// the chunk writer block until
	// all the chunks are received, reordered and written.
	d.chunkWriter(och)

	close(done)
	d.printProgress()

	return nil
}

func (d *download) showProgress(done <-chan struct{}) {
	d.printProgress()

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			d.printProgress()
		case <-done:
			return
		}
	}
}

// percentRatio returns the ration of two integers as a percentage.
func percentRatio(a, b int64) float64 {
	return float64(a) / float64(b) * 100.0
}

func (d *download) printProgress() {
	fmt.Fprintf(os.Stderr, "succeeded: %.2f%%; progress: %.2f%%; downloaded: %s; rate: %s/s\n", percentRatio(d.succeeded.value(), d.size), percentRatio(d.progress.value(), d.size), humanize.Bytes(uint64(d.downloaded.value())), humanize.Bytes(uint64(d.downloadMeter.Rate1())))
}

// contentSize returns the size of the resource
// at the given url.
func contentSize(url string) (int64, error) {
	resp, err := http.Get(url)
	if err != nil {
		return 0, err
	}
	return resp.ContentLength, nil
}

// followRedirect ensures that all headers are passed through the redirect.
// This is necessary for the range header to work across redirects.
func followRedirect(req *http.Request, via []*http.Request) error {
	req.Header = via[0].Header
	return nil
}

// fetch returns a ReadCloser with the content of a
// section of the resource found at url starting at from
// and len long.
func (d *download) fetch(from, len int64) (io.ReadCloser, error) {
	req, err := http.NewRequest("GET", d.url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", from, from+len-1))

	client := &http.Client{
		CheckRedirect: followRedirect,
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusPartialContent {
		defer resp.Body.Close()
		if resp.StatusCode == http.StatusOK {
			return nil, fmt.Errorf("Server doesn't implement partial content")
		}
		return nil, errors.New(resp.Status)
	}
	return &measuringReader{resp.Body, &d.downloaded, d.downloadMeter}, nil
}

type measuringReader struct {
	io.ReadCloser
	*counter
	metrics.Meter
}

func (m *measuringReader) Read(p []byte) (int, error) {
	n, err := m.ReadCloser.Read(p)
	m.add(int64(n))
	m.Mark(int64(n))
	return n, err
}

func main() {
	flag.Parse()
	if err := start(); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Fix resuming<commit_after>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/rcrowley/go-metrics"
)

var (
	output    = flag.String("o", "-", "output file")
	workers   = flag.Int("n", 4, "parallel connections")
	chunkSize = flag.Int("c", 100*1024, "chunk size")
	debug     = flag.Bool("d", false, "debug")
)

type chunk struct {
	offset int64
	size   int64
	file   string
}

type download struct {
	wg       sync.WaitGroup
	url      string
	size     int64
	startPos int64
	writer   io.Writer

	succeeded     counter
	downloaded    counter
	downloadMeter metrics.Meter
	progress      counter
}

type counter struct {
	v int64
}

func (c *counter) add(delta int64) {
	atomic.AddInt64(&c.v, delta)
}

func (c *counter) value() int64 {
	return atomic.LoadInt64(&c.v)
}

func newDownload(url string) (*download, error) {
	size, err := contentSize(url)
	if err != nil {
		return nil, err
	}

	return &download{
		url:           url,
		size:          size,
		writer:        os.Stdout,
		downloadMeter: metrics.NewMeter(),
	}, nil
}

// chunkFeeder generates chunk structures and sends them to the out channel.
// The chunk size is determined via commandline params. The last chunk is
// is sized accordingly.
func (d *download) chunkFeeder(out chan<- *chunk) {
	// when resuming a download. the chunk feeder won't emit chunks
	// that are already been written.
	d.succeeded.add(d.startPos)

	if *debug {
		log.Printf("Feeder is running")
	}
	for i := int64(0); ; i++ {
		size := int64(*chunkSize)
		off := i * size

		if off+size <= d.startPos {
			continue
		}
		if d.startPos-off > 0 && d.startPos-off < size {
			size -= d.startPos - off
			off = d.startPos
		}

		if off >= d.size {
			break
		}
		if off+size > d.size {
			size = off + size - d.size
		}
		if *debug {
			log.Printf("Enqueuing chunk %d size %d", off, size)
		}
		out <- &chunk{
			offset: off,
			size:   size}
	}
	if *debug {
		log.Printf("Feeder is closing")
	}
	close(out)
}

// chunkWorker task chunks from the in channel, fetches the chunk from
// the http resource into the buffer in the chunk structure.
// The completed chunk is then sent to the out channel.
// Any error causes to panic.
func (d *download) chunkWorker(n int, in <-chan *chunk, out chan<- *chunk) {
	if *debug {
		log.Println("Started chunk worker", n)
	}

	for ch := range in {
		if *debug {
			log.Printf("Chunk worker %d fetching chunk %d (%d)\n", n, ch.offset, ch.size)
		}

		basename := *output
		if *output == "-" {
			basename = "/tmp/pget.tmp"
		}

		offset := ch.offset
		size := ch.size

		var f *os.File
		var err error

		tmpName := fmt.Sprintf("%s.chunk-%d", basename, ch.offset)
		if st, err := os.Stat(tmpName); os.IsNotExist(err) {
			f, err = os.Create(tmpName)
			if err != nil {
				log.Fatalf("Cannot create temp file for chunk for offset %d, %s", ch.offset, err)
			}
		} else if err == nil {
			f, err = os.OpenFile(tmpName, os.O_APPEND|os.O_WRONLY, 0)
			if err != nil {
				log.Fatalf("Cannot open temp file for appending chunk for offset %d, %s", ch.offset, err)
			}
			offset += st.Size()
			size -= st.Size()
		} else {
			log.Fatal("Cannot stat ", tmpName, err)
		}

		var body io.ReadCloser
		if size < 1 {
			body = ioutil.NopCloser(&bytes.Buffer{})
		} else {
			if *debug {
				log.Printf("Worker %d fetching from %d -> size %d\n", n, offset, size)
			}

			body, err = d.fetch(offset, size)
			if err != nil {
				log.Fatal(err)
			}
			// mark the whole chunk as succeeded
			// even if we continued an interrupted download.
			d.succeeded.add(ch.size)
		}

		if _, err = io.Copy(f, body); err != nil {
			log.Fatal(err)
		}
		body.Close()
		f.Close()

		ch.file = tmpName
		out <- ch
	}

	if *debug {
		log.Println("Finishing chunk worker", n)
	}
}

// chunkWriter reads completed chunks from the in channel,
// and writes them out in the correct order.
// Out of order chunks are kept in memory until the missing pieces
// are delivered.
func (d *download) chunkWriter(in <-chan *chunk) {
	readyChunks := make(map[int64]*chunk)
	outPos := d.startPos
	d.progress.add(d.startPos)

	for ch := range in {
		if *debug {
			log.Printf("Chunk writer processing completed chunk %d (%d)\n", ch.offset, ch.size)
		}
		readyChunks[ch.offset] = ch

		for {
			if c, ok := readyChunks[outPos]; ok {
				f, err := os.Open(c.file)
				if err != nil {
					log.Fatal(err)
				}
				n, err := io.Copy(d.writer, f)
				if err != nil {
					log.Fatal(err)
				}
				d.progress.add(n)

				f.Close()
				//os.Remove(c.file)

				outPos += c.size
			} else {
				break
			}
		}
	}
}

// start
func start() error {
	url := flag.Arg(0)

	d, err := newDownload(url)
	if err != nil {
		return err
	}

	lastPos := int64(0)
	if *output != "-" {
		//f, err := os.Create(*output)
		f, err := os.OpenFile(*output, os.O_RDWR|os.O_CREATE, 0666)
		if err != nil {
			return err
		}
		defer f.Close()
		d.writer = f

		lastPos, err = f.Seek(0, 2)
		if err != nil {
			log.Fatal(err)
		}
	}
	d.startPos = lastPos

	if *debug {
		log.Println("Size", d.size)
		log.Println("Last pos", lastPos)
	}

	ich := make(chan *chunk, 2)
	och := make(chan *chunk, 2)

	// spawn the chunk feeder.
	go d.chunkFeeder(ich)

	// spawn the workers in such a way
	// that we know when they are all done.
	for i := 0; i < *workers; i++ {
		d.wg.Add(1)
		go func(i int) {
			d.chunkWorker(i, ich, och)
			d.wg.Done()
		}(i)
	}

	// when all the workers are gone
	// close the chunk writer input
	go func() {
		d.wg.Wait()
		close(och)
	}()

	done := make(chan struct{})
	go d.showProgress(done)

	// the chunk writer block until
	// all the chunks are received, reordered and written.
	d.chunkWriter(och)

	close(done)
	d.printProgress()

	return nil
}

func (d *download) showProgress(done <-chan struct{}) {
	d.printProgress()

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			d.printProgress()
		case <-done:
			return
		}
	}
}

// percentRatio returns the ration of two integers as a percentage.
func percentRatio(a, b int64) float64 {
	return float64(a) / float64(b) * 100.0
}

func (d *download) printProgress() {
	fmt.Fprintf(os.Stderr, "succeeded: %.2f%%; progress: %.2f%%; downloaded: %s; rate: %s/s\n", percentRatio(d.succeeded.value(), d.size), percentRatio(d.progress.value(), d.size), humanize.Bytes(uint64(d.downloaded.value())), humanize.Bytes(uint64(d.downloadMeter.Rate1())))
}

// contentSize returns the size of the resource
// at the given url.
func contentSize(url string) (int64, error) {
	resp, err := http.Get(url)
	if err != nil {
		return 0, err
	}
	return resp.ContentLength, nil
}

// followRedirect ensures that all headers are passed through the redirect.
// This is necessary for the range header to work across redirects.
func followRedirect(req *http.Request, via []*http.Request) error {
	req.Header = via[0].Header
	return nil
}

// fetch returns a ReadCloser with the content of a
// section of the resource found at url starting at from
// and len long.
func (d *download) fetch(from, len int64) (io.ReadCloser, error) {
	req, err := http.NewRequest("GET", d.url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", from, from+len-1))

	client := &http.Client{
		CheckRedirect: followRedirect,
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusPartialContent {
		defer resp.Body.Close()
		if resp.StatusCode == http.StatusOK {
			return nil, fmt.Errorf("Server doesn't implement partial content")
		}
		return nil, errors.New(resp.Status)
	}
	return &measuringReader{resp.Body, &d.downloaded, d.downloadMeter}, nil
}

type measuringReader struct {
	io.ReadCloser
	*counter
	metrics.Meter
}

func (m *measuringReader) Read(p []byte) (int, error) {
	n, err := m.ReadCloser.Read(p)
	m.add(int64(n))
	m.Mark(int64(n))
	return n, err
}

func main() {
	flag.Parse()
	if err := start(); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>