iplist/packed.go

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