m-dns/connection/connection.go

package connection

import (
	"net"
	"sync"
	"time"

	//"git.giaever.org/joachimmg/go-log.git/log"
	"git.giaever.org/joachimmg/m-dns/config"
)

type Conn interface {
	Read(b []byte) (int, *net.UDPAddr, error)
	Write(b []byte, addr *net.UDPAddr) (int, error)
	Close() error
	Lock()
	RLock()
	WLock()
	RUnlock()
	WUnlock()
	Unlock()
}

type UDP interface {
	Conn
	Listening() bool
	Listen(addr *net.UDPAddr) error
	ListenMulticast(iface *net.Interface, addr *net.UDPAddr) error
}

type conn struct {
	*net.UDPConn

	r sync.Mutex
	w sync.Mutex
}

type UDP4 struct {
	conn
}

type UDP6 struct {
	conn
}

func New(net int) UDP {
	if net == 4 {
		return new(UDP4)
	}
	return new(UDP6)
}

func (c *conn) Lock() {
	c.RLock()
	c.WLock()
}

func (c *conn) RLock() {
	c.r.Lock()
}

func (c *conn) WLock() {
	c.w.Lock()
}

func (c *conn) Unlock() {
	c.RUnlock()
	c.WUnlock()
}

func (c *conn) RUnlock() {
	c.r.Unlock()
}

func (c *conn) WUnlock() {
	c.w.Unlock()
}

func (c *conn) Listening() bool {
	return c.UDPConn != nil
}

func (c *conn) Read(b []byte) (int, *net.UDPAddr, error) {
	c.RLock()
	defer c.RUnlock()

	if !c.Listening() {
		return 0, nil, nil
	}

	if config.BufReadDeadline > 0 {
		c.SetReadDeadline(time.Now().Add(config.BufReadDeadline * time.Second))
	}

	return c.ReadFromUDP(b)
}

func (c *conn) Write(b []byte, addr *net.UDPAddr) (int, error) {
	c.WLock()
	defer c.WUnlock()

	if !c.Listening() {
		return 0, nil
	}

	if config.BufWriteDeadline > 0 {
		c.SetWriteDeadline(time.Now().Add(config.BufWriteDeadline * time.Second))
	}

	return c.WriteToUDP(b, addr)
}

func (c *conn) Close() error {
	c.Lock()
	defer c.Unlock()
	if !c.Listening() {
		return nil
	}

	err := c.UDPConn.Close()
	c.UDPConn = nil
	return err
}

func (u *UDP4) Listen(addr *net.UDPAddr) error {
	var err error
	u.UDPConn, err = net.ListenUDP("udp4", addr)
	return err
}

func (u *UDP4) ListenMulticast(iface *net.Interface, addr *net.UDPAddr) error {
	var err error
	u.UDPConn, err = net.ListenMulticastUDP("udp4", iface, addr)
	return err
}

func (u *UDP6) Listen(addr *net.UDPAddr) error {
	var err error
	u.UDPConn, err = net.ListenUDP("udp6", addr)
	return err
}

func (u *UDP6) ListenMulticast(iface *net.Interface, addr *net.UDPAddr) error {
	var err error
	u.UDPConn, err = net.ListenMulticastUDP("udp6", iface, addr)
	return err
}