iplist/packed.go

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