package haza

import "encoding/binary"
import "fmt"
import "net"
import "sync"

import "golang.org/x/net/bpf"
import "golang.org/x/sys/unix"

type Dhcp4UdpSock struct {
	addr *net.UDPAddr
	fd int
}

type Dhcp4Conn struct {
	fd int
	iface string

	udp_socks []*Dhcp4UdpSock
	mtx sync.Mutex
}

func NewDhcp4Conn(iface string, addr *net.UDPAddr) (*Dhcp4Conn, error) {
	var fd int
	var fd2 int
	var proto uint16
	var d *Dhcp4Conn
	var udp_socks []*Dhcp4UdpSock
	var ip net.IP
	var err error

	fd = -1
	fd2 = -1

	ip = addr.IP.To4()
	if ip == nil {
		err = fmt.Errorf("invalid address - %v", addr)
		goto oops
	}

	// Eventfd is linux specific. TODO: port it to other OSes.
	proto = Hton16(unix.ETH_P_IP)
	//proto = Hton16(unix.ETH_P_ALL)
	fd, err = unix.Socket(unix.AF_PACKET, unix.SOCK_RAW | unix.SOCK_CLOEXEC, int(proto))
	if err != nil { goto oops }

	fd2, err = open_udp4_socket(addr)
	if err != nil { goto oops }

	if iface != ""  {
		var sll *unix.SockaddrLinklayer
		var nif *net.Interface

		nif, err = net.InterfaceByName(iface)
		if err != nil { goto oops }

// TODO: if the interface is not an ethernet device, fail it.
/*
tun0
Interface type -> Encapsulation Type: Raw IP(7)

eth0
Encapsulation type: Ethernet(1)

unix.SockaddrLinklayer -> Hatype

/sys/class/net/tun0/type
 -> 65534 (ARPHRD_NONE)

/sys/class/net/eth0/type
 -> 1 (ARPHRD_ETHER)

#define ARPHRD_NETROM   0               * from KA9Q: NETROM pseudo    *
#define ARPHRD_ETHER    1               * Ethernet 10Mbps              *
#define ARPHRD_EETHER   2               * Experimental Ethernet        *
#define ARPHRD_AX25     3               * AX.25 Level 2                *
#define ARPHRD_PRONET   4               * PROnet token ring            *
#define ARPHRD_CHAOS    5               * Chaosnet                     *
#define ARPHRD_IEEE802  6               * IEEE 802.2 EthernetTRTB    *
#define ARPHRD_ARCNET   7               * ARCnet                       *
#define ARPHRD_APPLETLK 8               * APPLEtalk                    *
#define ARPHRD_DLCI     15              * Frame Relay DLCI             *
#define ARPHRD_ATM      19              * ATM                          *
#define ARPHRD_METRICOM 23              * Metricom STRIP (new IANA id) *
#define ARPHRD_IEEE1394 24              * IEEE 1394 IPv4 - RFC 2734    *
#define ARPHRD_EUI64    27              * EUI-64                       *
#define ARPHRD_INFINIBAND 32            * InfiniBand                   *

....
#define ARPHRD_VOID       0xFFFF        * Void type, nothing is known *
#define ARPHRD_NONE       0xFFFE        * zero header length *
*/

		sll = &unix.SockaddrLinklayer{
			Protocol: proto,
			Ifindex: nif.Index,
		}
		err = unix.Bind(fd, sll)
		if err != nil { goto oops }
	}

	udp_socks = make([]*Dhcp4UdpSock, 1)
	udp_socks[0] = &Dhcp4UdpSock{ fd: fd2, addr: addr }

	err = attach_dhcp4_filter(fd, iface, udp_socks)
	if err != nil { goto oops }

	d = &Dhcp4Conn{fd: fd, iface: iface, udp_socks: udp_socks}
	return d, nil

oops:
	if fd2 >= 0 { unix.Close(fd2) }
	if fd >= 0 { unix.Close(fd) }
	return nil, err
}

func (d *Dhcp4Conn) Close() {
	var us *Dhcp4UdpSock

	for _, us = range d.udp_socks {
		if us.fd >= 0 {
			unix.Close(us.fd)
			us.fd = -1
		}
	}

	if d.fd >= 0 {
		unix.Close(d.fd)
		d.fd = -1
	}
}

func (d *Dhcp4Conn) ForEachUdpSockFd(handler func (fd int)) {
	var us *Dhcp4UdpSock

	for _, us = range d.udp_socks {
		if us.fd >= 0 {
			// the handler must not close the fd
			handler(us.fd)
		}
	}
}

func (d *Dhcp4Conn) AddAddr(addr *net.UDPAddr) error {
	var fd int
	var udp_socks []*Dhcp4UdpSock
	var err error

	fd, err = open_udp4_socket(addr)
	if err != nil { return err }

	err = attach_discard_filter(fd)
	if err != nil { return err }

	d.mtx.Lock()
	defer d.mtx.Unlock()

	udp_socks = make([]*Dhcp4UdpSock, len(d.udp_socks) + 1)
	copy(udp_socks, d.udp_socks)
	udp_socks[len(d.udp_socks)] = &Dhcp4UdpSock{fd: fd, addr: addr}

	err = attach_dhcp4_filter(d.fd, d.iface, udp_socks)
	if err != nil {
		unix.Close(fd)
		return err
	}

	d.udp_socks = udp_socks
	return nil
}

/*
func (d *Dhcp4Conn) DelAddr(addr *net.UDPAddr) error {
	return nil
}*/

func sockaddrToString(sa unix.Sockaddr) string {
	switch v := sa.(type) {
	case *unix.SockaddrInet4:
		ip := net.IPv4(v.Addr[0], v.Addr[1], v.Addr[2], v.Addr[3])
		return fmt.Sprintf("%s:%d", ip.String(), v.Port)
	case *unix.SockaddrInet6:
		ip := net.IP(v.Addr[:])
		return fmt.Sprintf("[%s]:%d", ip.String(), v.Port)
	case *unix.SockaddrUnix:
		return v.Name
	case *unix.SockaddrLinklayer:
		return fmt.Sprintf("ifindex=%d, proto=0x%04x addr % x", v.Ifindex, Ntoh16(v.Protocol), v.Addr[:v.Halen])
	default:
		return fmt.Sprintf("unknown sockaddr type: %T", v)
	}
}

func patch_jump_if_skip_true(insts *[]bpf.Instruction, i uint8, offset uint8) {
	var inst bpf.JumpIf
	inst = (*insts)[i].(bpf.JumpIf)
	inst.SkipTrue = offset
	(*insts)[i] = inst
}

func patch_jump_if_skip_false(insts *[]bpf.Instruction, i uint8, offset uint8) {
	var inst bpf.JumpIf
	inst = (*insts)[i].(bpf.JumpIf)
	inst.SkipFalse = offset
	(*insts)[i] = inst
}

func attach_discard_filter(fd int) error {
	var insts []bpf.Instruction
	var raw_insts []bpf.RawInstruction
	var flts []unix.SockFilter
	var fprog *unix.SockFprog
	var i int
	var err error

	insts = []bpf.Instruction{
		bpf.RetConstant{Val: 0x0},
	}

	raw_insts, err = bpf.Assemble(insts)
	if err != nil { goto oops }

	flts = make([]unix.SockFilter, len(raw_insts))
	for i, _ = range raw_insts {
		flts[i] = unix.SockFilter{
			Code: raw_insts[i].Op,
			Jt: raw_insts[i].Jt,
			Jf: raw_insts[i].Jf,
			K: raw_insts[i].K,
		}
	}

	fprog = &unix.SockFprog{
		Len: uint16(len(flts)),
		Filter: &flts[0],
	}
	err = unix.SetsockoptSockFprog(fd, unix.SOL_SOCKET, unix.SO_ATTACH_FILTER, fprog)
	if err != nil { goto oops }

	return nil


oops:
	return err
}

func attach_dhcp4_filter(fd int, iface string, udp_socks []*Dhcp4UdpSock) error {
	var insts []bpf.Instruction
	var raw_insts []bpf.RawInstruction
	var flts []unix.SockFilter
	var fprog *unix.SockFprog
	var us *Dhcp4UdpSock
	var i int
	var ninsts int
	var ipos uint8
	var accept_pos uint8
	var reject_pos uint8
	var err error

	ninsts = 8;
	for _, us = range udp_socks {
		if us.addr.IP.To4() == nil { break }

		if !us.addr.IP.IsUnspecified() {
			ninsts = ninsts + 1
		}
		ninsts = ninsts + 3
	}

	ninsts = ninsts + 2
	if ninsts > 255 {
		if iface != "" {
			return fmt.Errorf("too many ip addresses on %s", iface)
		} else {
			return fmt.Errorf("too many ip addresses")
		}
	}

	insts = make([]bpf.Instruction, ninsts)

	// JumpIf instructions are patched below after ip address and port check instructions are produced
	insts[0] = bpf.LoadAbsolute{Off: 12, Size: 2} // load ethernet type(2 bytes) at offset 12.
	insts[1] = bpf.JumpIf{Cond: bpf.JumpEqual, Val: unix.ETH_P_IP, SkipTrue: 0, SkipFalse: 0}
	insts[2] = bpf.LoadAbsolute{Off: 14 + 9, Size: 1} // load IP protocol (offset 14 + 9)
	insts[3] = bpf.JumpIf{Cond: bpf.JumpEqual, Val: unix.IPPROTO_UDP, SkipTrue: 0, SkipFalse: 0}
	insts[4] = bpf.LoadAbsolute{Off: 14 + 6, Size: 2} // load IP flags & fragment offset (2 bytes)
	insts[5] = bpf.JumpIf{Cond: bpf.JumpBitsSet, Val: 0x1FFF, SkipTrue: 0, SkipFalse: 0}
	insts[6] = bpf.LoadAbsolute{Off: 14 + 16, Size: 4} // destination ip address
	insts[7] = bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0xFFFFFFFF, SkipTrue: 0, SkipFalse: 0}

	ipos = 8
	for _, us = range udp_socks {
		if us.addr.IP.To4() == nil { break }

		if !us.addr.IP.IsUnspecified() {
			var b net.IP
			// check more desination ip addresses
			b = us.addr.IP.To4()	 // must not fail as the check was done already
			insts[ipos] = bpf.JumpIf{Cond: bpf.JumpEqual, Val: binary.BigEndian.Uint32(b), SkipTrue: 0, SkipFalse: 3} // ip match. go to port match
			ipos = ipos + 1
		}

		// compute IP header length (IHL) into X register. LDX is like this:
		// +---------------+-----------------+---+---+---+
		// | AddrMode (3b) | LoadWidth (2b)  | 0 | 0 | 1 |
		// +---------------+-----------------+---+---+---+
		// this one is supposed to be LDX(0x1) | B(0x10) | MSH(0xA0)
		insts[ipos] = bpf.RawInstruction{Op: (0x1 | 0x10 | 0xA0), Jt: 0, Jf: 0, K: 14 }
		ipos = ipos + 1

		// load UDP destination port: offset = Ethernet(14) + 2 (UDP dest port) + IP header length * 4
		insts[ipos] = bpf.LoadIndirect{Off: 14 + 2, Size: 2}
		ipos = ipos + 1
		insts[ipos] = bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(us.addr.Port), SkipTrue: 0, SkipFalse: 0}
		ipos = ipos + 1
	}

	insts[ipos] = bpf.RetConstant{Val: 0} // reject
	reject_pos = ipos
	ipos = ipos + 1
	insts[ipos] = bpf.RetConstant{Val: 0xFFFF} // accept
	accept_pos = ipos
	ipos = ipos + 1

	patch_jump_if_skip_false(&insts, 1, reject_pos - 1 - 1) // discard if not ETH_P_IP
	patch_jump_if_skip_false(&insts, 3, reject_pos - 3 - 1) // discard if not UDP
	patch_jump_if_skip_true(&insts, 5, reject_pos - 5 - 1) // discard fragmented packets
	patch_jump_if_skip_true(&insts, 7, accept_pos - 7 - 1) // accept broadcast

	ipos = 8
	for _, us = range udp_socks {
		if us.addr.IP.To4() == nil { break }

		if !us.addr.IP.IsUnspecified() {
			ipos = ipos + 1
		}

		ipos = ipos + 1 // skip bpf.RawInstruction
		ipos = ipos + 1 // skip bpf.LoadIndirect
		patch_jump_if_skip_true(&insts, ipos, accept_pos - ipos - 1) // accept if port match
		ipos = ipos + 1
	}

	raw_insts, err = bpf.Assemble(insts)
	if err != nil { goto oops }

	flts = make([]unix.SockFilter, len(raw_insts))
	for i, _ = range raw_insts {
		flts[i] = unix.SockFilter{
			Code: raw_insts[i].Op,
			Jt: raw_insts[i].Jt,
			Jf: raw_insts[i].Jf,
			K: raw_insts[i].K,
		}
	}

	fprog = &unix.SockFprog{
		Len: uint16(len(flts)),
		Filter: &flts[0],
	}
	err = unix.SetsockoptSockFprog(fd, unix.SOL_SOCKET, unix.SO_ATTACH_FILTER, fprog)
	if err != nil { goto oops }

	return nil

oops:
	return err
}

func open_udp4_socket(addr *net.UDPAddr) (int, error) {
	var fd int
	var ip net.IP
	var err error

	ip = addr.IP.To4()
	if ip == nil {
		err = fmt.Errorf("invalid address - %v", addr)
		goto oops
	}

	// create a real udp packet to deceive the kernel.
	fd, err = unix.Socket(unix.AF_INET, unix.SOCK_DGRAM | unix.SOCK_CLOEXEC, unix.IPPROTO_UDP)
	if err != nil { goto oops }

	// we won't be using this socket to receive packets.
	// attach a filter that discard incoming packets
	err = attach_discard_filter(fd)
	if err != nil { goto oops }

	err = unix.Bind(fd, &unix.SockaddrInet4{ Addr: [4]byte(ip[0:4]), Port: addr.Port})
	if err != nil { goto oops }

	return fd, nil

oops:
	return -1, err
}