github.com/anacrolix/torrent@v1.61.0/iplist/packed.go (about)

     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  }