iplist/packed.go

   1 package iplist
   2
   3 import (
   4         "encoding/binary"
   5         "io"
   6         "net"
   7         "os"
   8
   9         "github.com/edsrzf/mmap-go"
  10 )
  11
  12 // The packed format is an 8 byte integer of the number of ranges. Then 20
  13 // bytes per range, consisting of 4 byte packed IP being the lower bound IP of
  14 // the range, then 4 bytes of the upper, inclusive bound, 8 bytes for the
  15 // offset of the description from the end of the packed ranges, and 4 bytes
  16 // for the length of the description. After these packed ranges, are the
  17 // concatenated descriptions.
  18
  19 const (
  20         packedRangesOffset = 8
  21         packedRangeLen     = 20
  22 )
  23
  24 func (me *IPList) WritePacked(w io.Writer) (err error) {
  25         descOffsets := make(map[string]int64, len(me.ranges))
  26         descs := make([]string, 0, len(me.ranges))
  27         var nextOffset int64
  28         // This is a little monadic, no?
  29         write := func(b []byte, expectedLen int) {
  30                 if err != nil {
  31                         return
  32                 }
  33                 var n int
  34                 n, err = w.Write(b)
  35                 if err != nil {
  36                         return
  37                 }
  38                 if n != expectedLen {
  39                         panic(n)
  40                 }
  41         }
  42         var b [8]byte
  43         binary.LittleEndian.PutUint64(b[:], uint64(len(me.ranges)))
  44         write(b[:], 8)
  45         for _, r := range me.ranges {
  46                 write(r.First.To4(), 4)
  47                 write(r.Last.To4(), 4)
  48                 descOff, ok := descOffsets[r.Description]
  49                 if !ok {
  50                         descOff = nextOffset
  51                         descOffsets[r.Description] = descOff
  52                         descs = append(descs, r.Description)
  53                         nextOffset += int64(len(r.Description))
  54                 }
  55                 binary.LittleEndian.PutUint64(b[:], uint64(descOff))
  56                 write(b[:], 8)
  57                 binary.LittleEndian.PutUint32(b[:], uint32(len(r.Description)))
  58                 write(b[:4], 4)
  59         }
  60         for _, d := range descs {
  61                 write([]byte(d), len(d))
  62         }
  63         return
  64 }
  65
  66 func NewFromPacked(b []byte) PackedIPList {
  67         return PackedIPList(b)
  68 }
  69
  70 type PackedIPList []byte
  71
  72 var _ Ranger = PackedIPList{}
  73
  74 func (me PackedIPList) len() int {
  75         return int(binary.LittleEndian.Uint64(me[:8]))
  76 }
  77
  78 func (me PackedIPList) NumRanges() int {
  79         return me.len()
  80 }
  81
  82 func (me PackedIPList) getFirst(i int) net.IP {
  83         off := packedRangesOffset + packedRangeLen*i
  84         return net.IP(me[off : off+4])
  85 }
  86
  87 func (me PackedIPList) getRange(i int) (ret Range) {
  88         rOff := packedRangesOffset + packedRangeLen*i
  89         last := me[rOff+4 : rOff+8]
  90         descOff := int(binary.LittleEndian.Uint64(me[rOff+8:]))
  91         descLen := int(binary.LittleEndian.Uint32(me[rOff+16:]))
  92         descOff += packedRangesOffset + packedRangeLen*me.len()
  93         ret = Range{
  94                 me.getFirst(i),
  95                 net.IP(last),
  96                 string(me[descOff : descOff+descLen]),
  97         }
  98         return
  99 }
 100
 101 func (me PackedIPList) Lookup(ip net.IP) (r Range, ok bool) {
 102         ip4 := ip.To4()
 103         if ip4 == nil {
 104                 // If the IP list was built successfully, then it only contained IPv4
 105                 // ranges. Therefore no IPv6 ranges are blocked.
 106                 if ip.To16() == nil {
 107                         r = Range{
 108                                 Description: "bad IP",
 109                         }
 110                         ok = true
 111                 }
 112                 return
 113         }
 114         return lookup(me.getFirst, me.getRange, me.len(), ip4)
 115 }
 116
 117 func MMapPacked(filename string) (ret Ranger, err error) {
 118         f, err := os.Open(filename)
 119         if os.IsNotExist(err) {
 120                 err = nil
 121                 return
 122         }
 123         if err != nil {
 124                 return
 125         }
 126         defer f.Close()
 127         mm, err := mmap.Map(f, mmap.RDONLY, 0)
 128         if err != nil {
 129                 return
 130         }
 131         ret = NewFromPacked(mm)
 132         return
 133 }