iplist/packed.go

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