From 533779c4a01c8fe814db9f5b0eb3bcc1ff3ec751 Mon Sep 17 00:00:00 2001 From: fu-l Date: Tue, 25 Feb 2020 10:15:08 +0800 Subject: [PATCH] add timeout --- cmd/root.go | 29 +++++++++++++---------------- cmd/search.go | 4 ++-- cmd/update.go | 3 +-- cmd/version.go | 2 +- main.go | 2 +- pkg/search/search.go | 12 ++++++++++-- pkg/search/search_test.go | 4 +--- pkg/update/update.go | 20 +++++++------------- pkg/update/update_test.go | 4 ---- pkg/utils/utils.go | 10 ++++------ pkg/utils/utils_test.go | 3 +-- 11 files changed, 41 insertions(+), 52 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index b2c04ad..5ae7da0 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -16,38 +16,36 @@ limitations under the License. package cmd import ( - "fmt" - "github.com/spf13/cobra" - "os" + "fmt" + "github.com/spf13/cobra" + "os" ) - //var cfgFile string - // rootCmd represents the base command when called without any subcommands var rootCmd = &cobra.Command{ - Use: "IP2Country", - Short: "Convert ip address to country code", - Long: `Convert ip address to country code. + Use: "IP2Country", + Short: "Convert ip address to country code", + Long: `Convert ip address to country code. Also it can update ipdata file from internet. For example: IP2Country update // update data file IP2Country search 8.8.8.8 // convert address to country code. IP2Country search 2c0f:ff10::12 // convert ipv6 address`, - // Uncomment the following line if your bare application - // has an action associated with it: - // Run: func(cmd *cobra.Command, args []string) { }, + // Uncomment the following line if your bare application + // has an action associated with it: + // Run: func(cmd *cobra.Command, args []string) { }, } // Execute adds all child commands to the root command and sets flags appropriately. // This is called by main.main(). It only needs to happen once to the rootCmd. func Execute() { - if err := rootCmd.Execute(); err != nil { - fmt.Println(err) - os.Exit(1) - } + if err := rootCmd.Execute(); err != nil { + fmt.Println(err) + os.Exit(1) + } } //func init() { @@ -91,4 +89,3 @@ func Execute() { // fmt.Println("Using config file:", viper.ConfigFileUsed()) // } //} - diff --git a/cmd/search.go b/cmd/search.go index 83405c3..1a29c39 100644 --- a/cmd/search.go +++ b/cmd/search.go @@ -26,8 +26,8 @@ import ( var searchCmd = &cobra.Command{ Use: "search", Short: "covert ip address to country code", - Long: `convert ip address to country code and accepted one ip address`, - Args: cobra.MinimumNArgs(1), + Long: `convert ip address to country code and accepted one ip address`, + Args: cobra.MinimumNArgs(1), Run: func(cmd *cobra.Command, args []string) { //fmt.Println("search called") fmt.Println(search.Do(strings.Join(args, " "))) diff --git a/cmd/update.go b/cmd/update.go index 755575d..138eb00 100644 --- a/cmd/update.go +++ b/cmd/update.go @@ -20,13 +20,12 @@ import ( "github.com/spf13/cobra" ) - //var force bool // updateCmd represents the update command var updateCmd = &cobra.Command{ Use: "update", Short: "update data file from internet", - Long: `update data file from internet and cached in $HOME/.IP2Country/`, + Long: `update data file from internet and cached in $HOME/.IP2Country/`, Run: func(cmd *cobra.Command, args []string) { //fmt.Println("update called") update.Do() diff --git a/cmd/version.go b/cmd/version.go index 4f773b6..1a98277 100644 --- a/cmd/version.go +++ b/cmd/version.go @@ -28,7 +28,7 @@ const ( var versionCmd = &cobra.Command{ Use: "version", Short: "show version of IP2Country", - Long: `show version of IP2Country`, + Long: `show version of IP2Country`, Run: func(cmd *cobra.Command, args []string) { fmt.Printf("Version: %s\n", Version) }, diff --git a/main.go b/main.go index 4acb089..8908e82 100644 --- a/main.go +++ b/main.go @@ -18,5 +18,5 @@ package main import "github.com/fumanne/IP2Country/cmd" func main() { - cmd.Execute() + cmd.Execute() } diff --git a/pkg/search/search.go b/pkg/search/search.go index b8f8431..e44cc1d 100644 --- a/pkg/search/search.go +++ b/pkg/search/search.go @@ -10,6 +10,7 @@ import ( "path/filepath" "regexp" "strings" + "time" ) func Do(ipAddress string) string { @@ -18,7 +19,7 @@ func Do(ipAddress string) string { os.Exit(0) } - if ! utils.IsIP(ipAddress) { + if !utils.IsIP(ipAddress) { fmt.Printf("Not A IPaddress Format: %s\n", ipAddress) os.Exit(0) } @@ -27,7 +28,14 @@ func Do(ipAddress string) string { for _, f := range setFiles(ipAddress) { go search(f, targetInt, ch) } - return <-ch + + select { + case data := <-ch: + return data + case <-time.After(time.Second * 15): + return "" + + } } func setFiles(IpAddress string) []string { diff --git a/pkg/search/search_test.go b/pkg/search/search_test.go index fa5ce5d..015a1a5 100644 --- a/pkg/search/search_test.go +++ b/pkg/search/search_test.go @@ -1,7 +1,6 @@ package search import ( - "testing" ) @@ -15,5 +14,4 @@ func TestDo(t *testing.T) { t.Errorf("Not Match") } - -} \ No newline at end of file +} diff --git a/pkg/update/update.go b/pkg/update/update.go index 8fb8288..0562ce7 100644 --- a/pkg/update/update.go +++ b/pkg/update/update.go @@ -42,7 +42,7 @@ func (r *Region) filename_v6() string { func (r *Region) stream() []byte { response, err := http.Get(r.url) - defer response.Body.Close() + defer func() { _ = response.Body.Close() }() utils.CheckErr(err) body, err := ioutil.ReadAll(response.Body) utils.CheckErr(err) @@ -53,12 +53,12 @@ func (r *Region) generate(wg *sync.WaitGroup) { defer wg.Done() f4, _ := os.OpenFile(r.file_v4(), os.O_RDWR|os.O_TRUNC|os.O_CREATE, os.ModePerm) f6, _ := os.OpenFile(r.file_v6(), os.O_RDWR|os.O_TRUNC|os.O_CREATE, os.ModePerm) - defer f4.Close() - defer f6.Close() - for _, record := range strings.Split(string(r.stream()), "\n") { - if ! isSkip(record) && isIPFlag(record) { + defer func() { _ = f4.Close() }() + defer func() { _ = f6.Close() }() + for _, record := range strings.Split(string(r.stream()), "\n") { + if !isSkip(record) && isIPFlag(record) { s, e, c := utils.ParseIPInt(record) - line := s.String() + "\t" + e.String() + "\t" + c + "\n" + line := s.String() + "\t" + e.String() + "\t" + c + "\n" if isV4record(record) { makeFile(f4, line) } @@ -75,8 +75,6 @@ func makeFile(file *os.File, s string) { utils.CheckErr(err) } - - func mkdir(d string) { if err := os.MkdirAll(d, os.ModePerm); err != nil { panic(err) @@ -90,7 +88,6 @@ func NewRegion(name, url string) *Region { } } - func isSkip(record string) bool { words := strings.Split(record, "|") if len(words) < 7 { @@ -106,14 +103,13 @@ func isIPFlag(record string) bool { if err != nil { panic(err) } - if ! ok { + if !ok { return false } return true } - func isV4record(record string) bool { words := strings.Split(record, "|") ok, _ := regexp.MatchString("^ipv4$", words[2]) @@ -134,8 +130,6 @@ func isV6record(record string) bool { } } - - func Do() { mkdir(utils.Locate(utils.DOWNLOAD)) wg := &sync.WaitGroup{} diff --git a/pkg/update/update_test.go b/pkg/update/update_test.go index 7456591..2159250 100644 --- a/pkg/update/update_test.go +++ b/pkg/update/update_test.go @@ -27,7 +27,6 @@ func TestIsSkip(t *testing.T) { } } - func TestIsIPFlag(t *testing.T) { correct := "afrinic|GA|ipv4|196.223.39.0|256|20140923|assigned" wrong := "afrinic|GA|Now|196.223.39.0|256|20140923|assigned" @@ -41,6 +40,3 @@ func TestIsIPFlag(t *testing.T) { } } - - - diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go index 2ec513e..1f8f868 100644 --- a/pkg/utils/utils.go +++ b/pkg/utils/utils.go @@ -16,9 +16,7 @@ const ( IPv6Prefix = "IPv6_" ) - -var PrivateIPRegexp = "^10.*$|^172.16.*$|^192.168.*$|^127.*$" - +var PrivateIPRegexp = "^10.*$|^172.(1[6-9]|2[0-9]|3[0-1]).*$|^192.168.*$|^127.*$" var IPv4Regexp = "^(([1-9]?[0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]).){3}([1-9]?[0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])$" var IPV6Regexp = `^((([0-9A-Fa-f]{1,4}:){7}([0-9A-Fa-f]{1,4}|:))|` + @@ -55,7 +53,7 @@ func Ip2long(ip string) *big.Int { func Str2BigInt(ipint string) *big.Int { x := big.NewInt(0) x, ok := x.SetString(ipint, 10) - if ! ok { + if !ok { panic("Set Ip to Big Int Error") } return x @@ -123,8 +121,8 @@ func IsPrivate(ip string) bool { } func IsIP(ip string) bool { - if ! IsIPv4(ip) && ! IsIPv6(ip) { + if !IsIPv4(ip) && !IsIPv6(ip) { return false } return true -} \ No newline at end of file +} diff --git a/pkg/utils/utils_test.go b/pkg/utils/utils_test.go index 5521a6a..2bb060e 100644 --- a/pkg/utils/utils_test.go +++ b/pkg/utils/utils_test.go @@ -14,7 +14,6 @@ func TestIP2long(t *testing.T) { } - func TestIsIPv4(t *testing.T) { right := "255.255.252.252" wrong := "fe80::21b:77ff:fbd6:7860" @@ -37,4 +36,4 @@ func TestIsIPv6(t *testing.T) { if IsIPv6(right) { t.Logf("%s is ipv6", right) } -} \ No newline at end of file +}