// Copyright 2022 The Ip2Region Authors. All rights reserved. // Use of this source code is governed by a Apache2.0-style // license that can be found in the LICENSE file. // original source ip editor package xdb import ( "container/list" "fmt" "io" "os" "path/filepath" "sort" ) type Editor struct { verison *Version // source ip file srcHandle io.ReadCloser toSave bool // segments list segments *list.List // region cache rgCache *RegionCache } func NewEditor(version *Version, srcFile string) (*Editor, error) { // check the src and dst file srcPath, err := filepath.Abs(srcFile) if err != nil { return nil, err } srcHandle, err := os.OpenFile(srcPath, os.O_RDONLY, 0600) if err != nil { return nil, err } return INewEditor(version, srcHandle) } func INewEditor(version *Version, srcReader io.ReadCloser) (*Editor, error) { e := &Editor{ verison: version, srcHandle: srcReader, toSave: false, segments: list.New(), rgCache: NewRegionCache(), } // load the segments if srcReader != nil { err := e.loadSegments() if err != nil { return nil, fmt.Errorf("failed to load segments: %s", err) } } return e, nil } func (e *Editor) Region(str string) *Region { return e.rgCache.Region(str) } // Load all the segments from the source file func (e *Editor) loadSegments() error { var last *Segment = nil var segments []*Segment var sorting = false _, _, iErr := IterateSegments(e.srcHandle, true, func(l string) { // do nothing here }, nil, e.Region, func(seg *Segment) error { // version check if len(seg.StartIP) != e.verison.Bytes { return fmt.Errorf("invalid ip segment(%s expected)", e.verison.Name) } // check the order of the data segment // if err := seg.RightBehind(last); err != nil { if err := seg.After(last); err != nil { // return err // @Note: If the continuity is disrupted, // we will sort all these segments later. sorting = true } // e.segments.PushBack(seg) segments = append(segments, seg) last = seg return nil }) if iErr != nil { return iErr } // check and do the sorting if sorting { sort.Slice(segments, func(i, j int) bool { return IPCompare(segments[i].StartIP, segments[j].StartIP) < 0 }) // open the to save e.toSave = true } // check and fill in the discontinuous segments // to Keep the entire data continuous. last = nil for _, seg := range segments { if last == nil { if IPCompare(seg.StartIP, e.verison.Min) > 0 { e.segments.PushBack(&Segment{ StartIP: e.verison.Min, EndIP: IPSubOne(seg.StartIP), Region: EmptyRegion, }) } } else if err := seg.RightBehind(last); err == nil { // Do nothing here since it just right behind the last } else if err := seg.After(last); err != nil { // segments overlap return fmt.Errorf("overlap checking: %w", err) } else { // push the padding segments e.segments.PushBack(&Segment{ StartIP: IPAddOne(last.EndIP), EndIP: IPSubOne(seg.StartIP), Region: EmptyRegion, }) } // push the current segment e.segments.PushBack(seg) // reset the last last = seg } // check and padding the tailing segmnet if back := e.segments.Back(); back != nil { if IPCompare(e.verison.Max, back.Value.(*Segment).EndIP) > 0 { e.segments.PushBack(&Segment{ StartIP: IPAddOne(back.Value.(*Segment).EndIP), EndIP: e.verison.Max, Region: EmptyRegion, }) } } segments = nil // let GC do it work return nil } func (e *Editor) NeedSave() bool { return e.toSave } func (e *Editor) SegLen() int { return e.segments.Len() } func (e *Editor) Slice(offset int, size int) []*Segment { var index = -1 var out []*Segment var next *list.Element for ele := e.segments.Front(); ele != nil; ele = next { next = ele.Next() s, ok := ele.Value.(*Segment) if !ok { continue } // offset match index++ if index < offset { continue } out = append(out, s) if len(out) >= size { break } } return out } func (e *Editor) Put(ip string, cb func(newSeg *Segment, oldList []*Segment) []*Segment) (int, int, error) { seg, err := SegmentFrom(ip, e.rgCache.Region) if err != nil { return 0, 0, err } return e.PutSegment(seg, cb) } // PutSegment put the specified segment into the current segment list with // the following position relationships. // 1, A - fully contained like: // StartIP------seg.StartIP--------seg.EndIP----EndIP // // |------------------| // // 2, B - intersect like: // StartIP------seg.StartIP------EndIP------| // // |---------------------seg.EndIP func (e *Editor) PutSegment(seg *Segment, cb func(newSeg *Segment, oldList []*Segment) []*Segment) (int, int, error) { var next *list.Element var eList []*list.Element var found, counter = false, 0 for ele := e.segments.Front(); ele != nil; ele = next { next = ele.Next() s, ok := ele.Value.(*Segment) if !ok { // could this even be a case ? return 0, 0, fmt.Errorf("type error: ele not a Segment ptr") } counter++ // found the related segment if found { // just keep going } else if IPCompare(seg.StartIP, s.StartIP) >= 0 && IPCompare(seg.StartIP, s.EndIP) <= 0 { found = true } else { continue } eList = append(eList, ele) if IPCompare(seg.EndIP, s.EndIP) <= 0 { break } } if len(eList) == 0 { // could this even be a case ? // return 0, 0, fmt.Errorf("failed to find the related segment") e.segments.PushBack(seg) return 0, 1, nil } // print for debug // for i, s := range eList { // fmt.Printf("find -> %d: %s\n", i, s.Value.(*Segment)) // } // segment split var sList []*Segment var head = eList[0].Value.(*Segment) if IPCompare(seg.StartIP, head.StartIP) > 0 { sList = append(sList, &Segment{ StartIP: head.StartIP, EndIP: IPSubOne(seg.StartIP), Region: head.Region, }) } // check the callback and append the new segment if cb == nil { sList = append(sList, seg) } else { // call the callback and append the segments sList = append(sList, cb(seg, e._getSegments(seg, eList))...) } // check and do the tailing segment append var tail = eList[len(eList)-1].Value.(*Segment) if IPCompare(seg.EndIP, tail.EndIP) < 0 { sList = append(sList, &Segment{ StartIP: IPAddOne(seg.EndIP), EndIP: tail.EndIP, Region: tail.Region, }) } // print for debug // for i, s := range sList { // fmt.Printf("replace -> %d: %s\n", i, s) // } // delete all the in-range segments and var base *list.Element var oldRows, newRows = len(eList), len(sList) for _, ele := range eList { base = ele.Next() e.segments.Remove(ele) } // add all the new segments if base == nil { for _, s := range sList { e.segments.PushBack(s) } } else { for _, s := range sList { e.segments.InsertBefore(s, base) } } // open the to save flag e.toSave = true return oldRows, newRows, nil } func (e *Editor) _getSegments(seg *Segment, eList []*list.Element) []*Segment { var sList []*Segment var eLen = len(eList) switch eLen { case 0: return sList case 1: fSeg := eList[0].Value.(*Segment) return append(sList, &Segment{ StartIP: seg.StartIP, EndIP: seg.EndIP, Region: fSeg.Region, }) } var maxIdx = eLen - 1 // append the first segment fSeg := eList[0].Value.(*Segment) sList = append(sList, &Segment{ StartIP: seg.StartIP, EndIP: fSeg.EndIP, Region: fSeg.Region, }) // append the middle segments for i := 1; i < maxIdx; i++ { sList = append(sList, eList[i].Value.(*Segment)) } // append the last segment lSeg := eList[maxIdx].Value.(*Segment) sList = append(sList, &Segment{ StartIP: lSeg.StartIP, EndIP: seg.EndIP, Region: lSeg.Region, }) return sList } func (e *Editor) PutFile(src string, cb func(newSeg *Segment, oldList []*Segment) []*Segment) (int, int, error) { handle, err := os.OpenFile(src, os.O_RDONLY, 0600) if err != nil { return 0, 0, err } var oldRows, newRows = 0, 0 _, _, iErr := IterateSegments(handle, true, func(l string) { // do nothing here }, nil, NewRegion, func(seg *Segment) error { o, n, err := e.PutSegment(seg, cb) if err == nil { oldRows += o newRows += n } return err }) if iErr != nil { return oldRows, newRows, iErr } _ = handle.Close() return oldRows, newRows, nil } func (e *Editor) SaveToFile(dstFile string) error { // check the to-save flag if !e.toSave { return fmt.Errorf("nothing changed") } dstHandle, err := os.OpenFile(dstFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644) if err != nil { return err } defer dstHandle.Close() // loop and flush all the segments to the dstHandle var next *list.Element for ele := e.segments.Front(); ele != nil; ele = next { next = ele.Next() s, ok := ele.Value.(*Segment) if !ok { // could this even be a case ? continue } // ignore the padded or empty segment if s.Region.IsEmpty() { continue } // var l = s.String() // _, err = dstHandle.WriteString(fmt.Sprintf("%s\n", l)) _, err = fmt.Fprintln(dstHandle, s.String()) if err != nil { return err } } // close the to-save flag e.toSave = false return nil } func (e *Editor) Close() { _ = e.srcHandle.Close() e.rgCache.Clean() }