CIDR tree for blacklist matching purpose.

cidr_tree.go

package utils

/* 
 * Charles-Antoine.Mathieu <skatkatt@root.gg>
 * CIDR tree for blacklist matching purpose.
 *
 * It's actually a simple binary tree with each node matching a byte
 * of the cidr until the mask is exhausted. The tree has to be walked
 * from the root following one of the two available pointers depending
 * if the ip to be matched is above of below the median stored in the node.
 * So it will take maximum 32 call to match an ip against a /32.
 *
 * Any data can be attached to a node but at least something different than
 * nil must be attached to differentiate empty nodes ( just here to link
 * with higher levels ).
 */

import (
	"net"
	"errors"
	"log"
)

/*
 * This is a node in the tree, it gets 2 pointer
 */
type CidrNode struct {
	Cidr     	*CIDR
	Median     	uint32
	Data       	interface{}
	NextLeft    *CidrNode
	NextRight   *CidrNode
}

/*
 * Constructor
 */
func NewCidrNode(cidr *CIDR, data interface {}) (this *CidrNode) {
	this = new(CidrNode)
	this.Cidr = cidr
	this.Data = data

	if (this.Cidr.Size == 1) {
		// /32
		this.Median = 0
	} else {
		this.Median = this.Cidr.First + uint32( ( this.Cidr.Size + 1 ) / 2 )
	}
	return
}

/*
 * Constructor using string representation of the CIDR.
 */
func NewCidrNodeFromString(str string, data interface {}) (*CidrNode, error) {
	if cidr, err := NewCIDRFromString(str, false) ; err == nil {
		return NewCidrNode(cidr,data), nil
	} else {
		return nil, err
	}
}

/*
 * Insert a higher node at the left or at the right.
 */
func (this *CidrNode) SetNext(node *CidrNode) {
	if (node.Cidr.First >= this.Median) {
		if (this.NextRight != nil) {
			this.NextRight.Data = node.Data
		} else {
			this.NextRight = node
		}
	} else {
		if (this.NextLeft != nil) {
			this.NextLeft.Data = node.Data
		} else {
			this.NextLeft = node
		}
	}
}

/*
 * Return the next node.
 */
func (this *CidrNode) GetNext(ip uint32) *CidrNode{
	if (ip >= this.Median) {
		return this.NextRight
	} else {
		return this.NextLeft
	}
}

/*
 * Display the CIDR of the node
 * TODO : display the data too
 */
func (this *CidrNode) String() string {
	return this.Cidr.String()
}

/*
 * This is the actual tree. It has to be walked
 * from the root to the last matching node.
 */
type CidrTree struct {
	root *CidrNode
}

/*
 * Constructor
 * The root of the tree is at 0.0.0.0/0
 */
func NewCidrTree(data interface {}) (this *CidrTree) {
	this = new(CidrTree)
	var err error
	if this.root, err = NewCidrNodeFromString("0.0.0.0/0", data) ; err != nil {
		log.Fatal(err)
	}
	return
}

/*
 * Add a cidr to the tree.
 */
func (this *CidrTree) AddNode(node *CidrNode) {
	parent := this.root
	//fmt.Println("add " + node.Cidr.String())

	// Skip all already inserted nodes
	for {
		if next := parent.GetNext(node.Cidr.First) ; next != nil {
			if next.Cidr.Slash < node.Cidr.Slash {
				//fmt.Println("skip " + next.String())
				parent = next
			} else {
				return
			}
		} else {
			break
		}
	}

	// Insert all missing empty nodes
	for parent.Cidr.Slash != node.Cidr.Slash - 1 {
		nextCidr := parent.Cidr

		cidr := new(net.IPNet)

		// Increment mask ( add one true bit at the end )
		cidr.Mask = net.IPMask(ItoN(NtoI(net.IP(parent.Cidr.Cidr.Mask)) >> 1 | 0x80000000))

		// The new CIDR is either the firt or the last half of the parent CIDR
		if (node.Cidr.First >= parent.Median){
			cidr.IP = ItoN(parent.Median)
		} else {
			cidr.IP = parent.Cidr.Cidr.IP
		}

		nextCidr, _ = NewCIDR(cidr,false)
		tmpNode := NewCidrNode(nextCidr, nil)

		//fmt.Println("add tmp node  " + tmpNode.String())
		parent.SetNext(tmpNode)
		parent = tmpNode
	}

	// Insert the real node
	//fmt.Println("add node  " + node.String())
	parent.SetNext(node)
}

func (this *CidrTree) AddNodeFromString(str string, data interface {}) error {
	if node, err := NewCidrNodeFromString(str,data) ; err == nil {
		this.AddNode(node)
		return nil
	} else {
		return err
	}
}

/*
 * Look for a node matching the IP ( a CIDR embracing this IP
 * has been added to the tree before ).
 */
func (this *CidrTree) Match(ip uint32) (node *CidrNode) {
	current := this.root
	for {
		if next := current.GetNext(ip) ; next != nil {
			current = next
			if ( current.Data != nil ) {
				//fmt.Println("found " + current.String())
				node = current
			} else {
				//fmt.Println("matched " + current.String())
			}
		} else {
			break
		}
	}
	return
}

/*
 * Look for a node matching the IP ( a CIDR embracing this IP
 * has been added to the tree before ).
 */
func (this *CidrTree) MatchFromString(str string) (*CidrNode, error) {
	if ip := net.ParseIP(str).To4() ; ip != nil {
		//fmt.Printf("matching %s %d\n",str, NtoI(ip))
		return this.Match(NtoI(ip)), nil
	} else {
		return nil, errors.New("Invalid IP : " + str)
	}
}