torrent/iplist/packed.go

146 lines
3.1 KiB
Go

//go:build !wasm
// +build !wasm
package iplist
import (
"encoding/binary"
"fmt"
"io"
"net"
"os"
"github.com/edsrzf/mmap-go"
)
// The packed format is an 8 byte integer of the number of ranges. Then 20
// bytes per range, consisting of 4 byte packed IP being the lower bound IP of
// the range, then 4 bytes of the upper, inclusive bound, 8 bytes for the
// offset of the description from the end of the packed ranges, and 4 bytes
// for the length of the description. After these packed ranges, are the
// concatenated descriptions.
const (
packedRangesOffset = 8
packedRangeLen = 44
)
func (ipl *IPList) WritePacked(w io.Writer) (err error) {
descOffsets := make(map[string]int64, len(ipl.ranges))
descs := make([]string, 0, len(ipl.ranges))
var nextOffset int64
// This is a little monadic, no?
write := func(b []byte, expectedLen int) {
if err != nil {
return
}
var n int
n, err = w.Write(b)
if err != nil {
return
}
if n != expectedLen {
panic(n)
}
}
var b [8]byte
binary.LittleEndian.PutUint64(b[:], uint64(len(ipl.ranges)))
write(b[:], 8)
for _, r := range ipl.ranges {
write(r.First.To16(), 16)
write(r.Last.To16(), 16)
descOff, ok := descOffsets[r.Description]
if !ok {
descOff = nextOffset
descOffsets[r.Description] = descOff
descs = append(descs, r.Description)
nextOffset += int64(len(r.Description))
}
binary.LittleEndian.PutUint64(b[:], uint64(descOff))
write(b[:], 8)
binary.LittleEndian.PutUint32(b[:], uint32(len(r.Description)))
write(b[:4], 4)
}
for _, d := range descs {
write([]byte(d), len(d))
}
return
}
func NewFromPacked(b []byte) PackedIPList {
ret := PackedIPList(b)
minLen := packedRangesOffset + ret.len()*packedRangeLen
if len(b) < minLen {
panic(fmt.Sprintf("packed len %d < %d", len(b), minLen))
}
return ret
}
type PackedIPList []byte
var _ Ranger = PackedIPList{}
func (pil PackedIPList) len() int {
return int(binary.LittleEndian.Uint64(pil[:8]))
}
func (pil PackedIPList) NumRanges() int {
return pil.len()
}
func (pil PackedIPList) getFirst(i int) net.IP {
off := packedRangesOffset + packedRangeLen*i
return net.IP(pil[off : off+16])
}
func (pil PackedIPList) getRange(i int) (ret Range) {
rOff := packedRangesOffset + packedRangeLen*i
last := pil[rOff+16 : rOff+32]
descOff := int(binary.LittleEndian.Uint64(pil[rOff+32:]))
descLen := int(binary.LittleEndian.Uint32(pil[rOff+40:]))
descOff += packedRangesOffset + packedRangeLen*pil.len()
ret = Range{
pil.getFirst(i),
net.IP(last),
string(pil[descOff : descOff+descLen]),
}
return
}
func (pil PackedIPList) Lookup(ip net.IP) (r Range, ok bool) {
ip16 := ip.To16()
if ip16 == nil {
panic(ip)
}
return lookup(pil.getFirst, pil.getRange, pil.len(), ip16)
}
type closerFunc func() error
func (me closerFunc) Close() error {
return me()
}
func MMapPackedFile(filename string) (
ret interface {
Ranger
io.Closer
},
err error,
) {
f, err := os.Open(filename)
if err != nil {
return
}
defer f.Close()
mm, err := mmap.Map(f, mmap.RDONLY, 0)
if err != nil {
return
}
ret = struct {
Ranger
io.Closer
}{NewFromPacked(mm), closerFunc(mm.Unmap)}
return
}