github.com/infraboard/keyauth@v0.8.1/apps/ip2region/impl/ip2region.go (about)

     1  package impl
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"os"
     7  
     8  	"github.com/infraboard/mcube/exception"
     9  
    10  	"github.com/infraboard/keyauth/apps/ip2region"
    11  	"github.com/infraboard/keyauth/apps/ip2region/reader"
    12  	"github.com/infraboard/keyauth/apps/storage"
    13  )
    14  
    15  func (s *service) UpdateDBFile(req *ip2region.UpdateDBFileRequest) error {
    16  	if err := req.Validate(); err != nil {
    17  		return exception.NewBadRequest("validate update db file requrest error, %s", err)
    18  	}
    19  
    20  	reader := req.ReadCloser()
    21  	defer reader.Close()
    22  
    23  	uploadReq := storage.NewUploadFileRequest(s.bucketName, s.dbFileName, reader)
    24  	uploadReq.WithToken(req.GetToken())
    25  	return s.storage.UploadFile(uploadReq)
    26  }
    27  
    28  func (s *service) LookupIP(ip string) (*ip2region.IPInfo, error) {
    29  	dbReader, err := s.getDBReader()
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  
    34  	return dbReader.MemorySearch(ip)
    35  }
    36  
    37  func (s *service) getDBReader() (*reader.IPReader, error) {
    38  	s.Lock()
    39  	defer s.Unlock()
    40  
    41  	if s.dbReader != nil {
    42  		return s.dbReader, nil
    43  	}
    44  
    45  	// 优先从本地文件加载DB文件
    46  	if err := s.loadDBFileFromLocal(); err != nil {
    47  		s.log.Infof("load ip2region db file from local error, %s, retry other load method ", err)
    48  	} else {
    49  		return s.dbReader, nil
    50  	}
    51  
    52  	if err := s.loadDBFileFromBucket(); err != nil {
    53  		s.log.Infof("load ip2region db file from bucket error, %s", err)
    54  	} else {
    55  		return s.dbReader, nil
    56  	}
    57  
    58  	return nil, fmt.Errorf("load ip2region db file error")
    59  }
    60  
    61  func (s *service) loadDBFileFromLocal() error {
    62  	file, err := os.Open(s.dbFileName)
    63  	if err != nil {
    64  		return fmt.Errorf("open file error, %s", err)
    65  	}
    66  
    67  	reader, err := reader.New(file)
    68  	if err != nil {
    69  		return err
    70  	}
    71  	s.dbReader = reader
    72  	return nil
    73  }
    74  
    75  func (s *service) loadDBFileFromBucket() error {
    76  	buf := bytes.NewBuffer([]byte{})
    77  	downloadReq := storage.NewDownloadFileRequest(s.bucketName, s.dbFileName, buf)
    78  	if err := s.storage.Download(downloadReq); err != nil {
    79  		return err
    80  	}
    81  
    82  	reader, err := reader.New(buf)
    83  	if err != nil {
    84  		return err
    85  	}
    86  	s.dbReader = reader
    87  
    88  	return nil
    89  }