feat: add --ipv4 & --ipv6 flags (jsdelivr#113)

radulucut · Jul 7, 2024 · 003bfb9 · 003bfb9
1 parent 03a8dfa
commit 003bfb9
Show file tree

Hide file tree

Showing 16 changed files with 502 additions and 16 deletions.
diff --git a/README.md b/README.md
@@ -105,6 +105,8 @@ Flags:
                        - an ID of a previous measurement to run with its probes
                        (default "world")
   -h, --help          help for globalping
+  -4, --ipv4          resolve names to IPv4 addresses
+  -6, --ipv6          resolve names to IPv6 addresses
   -J, --json          output results in JSON format (default false)
       --latency       output only the latency stats; applicable only to dns, http, and ping
                       commands (default false)

diff --git a/cmd/common.go b/cmd/common.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
+	"net"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -19,9 +20,11 @@ import (
 )
 
 var (
-	ErrNoPreviousMeasurements = errors.New("no previous measurements found")
-	ErrInvalidIndex           = errors.New("invalid index")
-	ErrIndexOutOfRange        = errors.New("index out of range")
+	ErrNoPreviousMeasurements      = errors.New("no previous measurements found")
+	ErrInvalidIndex                = errors.New("invalid index")
+	ErrIndexOutOfRange             = errors.New("index out of range")
+	ErrTargetIPVersionNotAllowed   = errors.New("ipVersion is not allowed when target is not a domain")
+	ErrResolverIPVersionNotAllowed = errors.New("ipVersion is not allowed when resolver is not a domain")
 )
 var (
 	saveIdToSessionErr = "failed to save measurement ID: %s"
@@ -48,6 +51,15 @@ func (r *Root) updateContext(cmd string, args []string) error {
 		r.ctx.Resolver = targetQuery.Resolver
 	}
 
+	if r.ctx.Ipv4 || r.ctx.Ipv6 {
+		if net.ParseIP(r.ctx.Target) != nil {
+			return ErrTargetIPVersionNotAllowed
+		}
+		if r.ctx.Resolver != "" && net.ParseIP(r.ctx.Resolver) != nil {
+			return ErrResolverIPVersionNotAllowed
+		}
+	}
+
 	// Check env for CI
 	if os.Getenv("CI") != "" {
 		r.ctx.CIMode = true

diff --git a/cmd/common_test.go b/cmd/common_test.go
@@ -10,11 +10,13 @@ import (
 
 func Test_UpdateContext(t *testing.T) {
 	for scenario, fn := range map[string]func(t *testing.T){
-		"no_arg":             test_updateContext_NoArg,
-		"country":            test_updateContext_Country,
-		"country_whitespace": test_updateContext_CountryWhitespace,
-		"no_target":          test_updateContext_NoTarget,
-		"ci_env":             test_uodateContext_CIEnv,
+		"no_arg":                test_updateContext_NoArg,
+		"country":               test_updateContext_Country,
+		"country_whitespace":    test_updateContext_CountryWhitespace,
+		"no_target":             test_updateContext_NoTarget,
+		"ci_env":                test_updateContext_CIEnv,
+		"target_not_hostname":   test_updateContext_TargetIsNotAHostname,
+		"resolver_not_hostname": test_updateContext_ResolverIsNotAHostname,
 	} {
 		t.Run(scenario, func(t *testing.T) {
 			fn(t)
@@ -68,7 +70,7 @@ func test_updateContext_NoTarget(t *testing.T) {
 	assert.Error(t, err)
 }
 
-func test_uodateContext_CIEnv(t *testing.T) {
+func test_updateContext_CIEnv(t *testing.T) {
 	oldCI := os.Getenv("CI")
 	t.Setenv("CI", "true")
 	defer t.Setenv("CI", oldCI)
@@ -85,6 +87,36 @@ func test_uodateContext_CIEnv(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+func test_updateContext_TargetIsNotAHostname(t *testing.T) {
+	ctx := createDefaultContext("ping")
+	ctx.Ipv4 = true
+	printer := view.NewPrinter(nil, nil, nil)
+	root := NewRoot(printer, ctx, nil, nil, nil, nil)
+
+	err := root.updateContext("ping", []string{"1.1.1.1"})
+	assert.EqualError(t, err, ErrTargetIPVersionNotAllowed.Error())
+
+	ctx.Ipv4 = false
+	ctx.Ipv6 = true
+	err = root.updateContext("ping", []string{"1.1.1.1"})
+	assert.EqualError(t, err, ErrTargetIPVersionNotAllowed.Error())
+}
+
+func test_updateContext_ResolverIsNotAHostname(t *testing.T) {
+	ctx := createDefaultContext("dns")
+	ctx.Ipv4 = true
+	printer := view.NewPrinter(nil, nil, nil)
+	root := NewRoot(printer, ctx, nil, nil, nil, nil)
+
+	err := root.updateContext("dns", []string{"example.com", "@1.1.1.1"})
+	assert.EqualError(t, err, ErrResolverIPVersionNotAllowed.Error())
+
+	ctx.Ipv4 = false
+	ctx.Ipv6 = true
+	err = root.updateContext("dns", []string{"example.com", "@1.1.1.1"})
+	assert.EqualError(t, err, ErrResolverIPVersionNotAllowed.Error())
+}
+
 func Test_ParseTargetQuery_Simple(t *testing.T) {
 	cmd := "ping"
 	args := []string{"example.com"}

diff --git a/cmd/dns.go b/cmd/dns.go
@@ -93,6 +93,12 @@ func (r *Root) RunDNS(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	if r.ctx.Ipv4 {
+		opts.Options.IPVersion = globalping.IPVersion4
+	} else if r.ctx.Ipv6 {
+		opts.Options.IPVersion = globalping.IPVersion6
+	}
+
 	res, showHelp, err := r.client.CreateMeasurement(opts)
 	if err != nil {
 		if !showHelp {

diff --git a/cmd/dns_test.go b/cmd/dns_test.go
@@ -82,3 +82,83 @@ func Test_Execute_DNS_Default(t *testing.T) {
 	)
 	assert.Equal(t, expectedHistory, string(b))
 }
+
+func Test_Execute_DNS_IPv4(t *testing.T) {
+	t.Cleanup(sessionCleanup)
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	expectedOpts := createDefaultMeasurementCreate("dns")
+	expectedOpts.Options.IPVersion = globalping.IPVersion4
+	expectedOpts.Options.Query = &globalping.QueryOptions{}
+
+	expectedResponse := createDefaultMeasurementCreateResponse()
+
+	gbMock := mocks.NewMockClient(ctrl)
+	gbMock.EXPECT().CreateMeasurement(expectedOpts).Times(1).Return(expectedResponse, false, nil)
+
+	viewerMock := mocks.NewMockViewer(ctrl)
+	viewerMock.EXPECT().Output(measurementID1, expectedOpts).Times(1).Return(nil)
+
+	timeMock := mocks.NewMockTime(ctrl)
+	timeMock.EXPECT().Now().Return(defaultCurrentTime).AnyTimes()
+
+	w := new(bytes.Buffer)
+	printer := view.NewPrinter(nil, w, w)
+	ctx := createDefaultContext("dns")
+	root := NewRoot(printer, ctx, viewerMock, timeMock, gbMock, nil)
+
+	os.Args = []string{"globalping", "dns", "jsdelivr.com",
+		"from", "Berlin",
+		"--ipv4"}
+	err := root.Cmd.ExecuteContext(context.TODO())
+	assert.NoError(t, err)
+
+	assert.Equal(t, "", w.String())
+
+	expectedCtx := createDefaultExpectedContext("dns")
+	expectedCtx.Ipv4 = true
+
+	assert.Equal(t, expectedCtx, ctx)
+}
+
+func Test_Execute_DNS_IPv6(t *testing.T) {
+	t.Cleanup(sessionCleanup)
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	expectedOpts := createDefaultMeasurementCreate("dns")
+	expectedOpts.Options.IPVersion = globalping.IPVersion6
+	expectedOpts.Options.Query = &globalping.QueryOptions{}
+
+	expectedResponse := createDefaultMeasurementCreateResponse()
+
+	gbMock := mocks.NewMockClient(ctrl)
+	gbMock.EXPECT().CreateMeasurement(expectedOpts).Times(1).Return(expectedResponse, false, nil)
+
+	viewerMock := mocks.NewMockViewer(ctrl)
+	viewerMock.EXPECT().Output(measurementID1, expectedOpts).Times(1).Return(nil)
+
+	timeMock := mocks.NewMockTime(ctrl)
+	timeMock.EXPECT().Now().Return(defaultCurrentTime).AnyTimes()
+
+	w := new(bytes.Buffer)
+	printer := view.NewPrinter(nil, w, w)
+	ctx := createDefaultContext("dns")
+	root := NewRoot(printer, ctx, viewerMock, timeMock, gbMock, nil)
+
+	os.Args = []string{"globalping", "dns", "jsdelivr.com",
+		"from", "Berlin",
+		"--ipv6"}
+	err := root.Cmd.ExecuteContext(context.TODO())
+	assert.NoError(t, err)
+
+	assert.Equal(t, "", w.String())
+
+	expectedCtx := createDefaultExpectedContext("dns")
+	expectedCtx.Ipv6 = true
+
+	assert.Equal(t, expectedCtx, ctx)
+}
diff --git a/cmd/http.go b/cmd/http.go
@@ -102,6 +102,12 @@ func (r *Root) RunHTTP(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	if r.ctx.Ipv4 {
+		opts.Options.IPVersion = globalping.IPVersion4
+	} else if r.ctx.Ipv6 {
+		opts.Options.IPVersion = globalping.IPVersion6
+	}
+
 	res, showHelp, err := r.client.CreateMeasurement(opts)
 	if err != nil {
 		if !showHelp {

diff --git a/cmd/http_test.go b/cmd/http_test.go
@@ -91,6 +91,94 @@ func Test_Execute_HTTP_Default(t *testing.T) {
 	assert.Equal(t, expectedHistory, string(b))
 }
 
+func Test_Execute_HTTP_IPv4(t *testing.T) {
+	t.Cleanup(sessionCleanup)
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	expectedOpts := createDefaultMeasurementCreate("http")
+	expectedOpts.Options.Protocol = "http"
+	expectedOpts.Options.IPVersion = globalping.IPVersion4
+	expectedOpts.Options.Request = &globalping.RequestOptions{
+		Headers: map[string]string{},
+		Host:    "jsdelivr.com",
+	}
+
+	expectedResponse := createDefaultMeasurementCreateResponse()
+
+	gbMock := mocks.NewMockClient(ctrl)
+	gbMock.EXPECT().CreateMeasurement(expectedOpts).Times(1).Return(expectedResponse, false, nil)
+
+	viewerMock := mocks.NewMockViewer(ctrl)
+	viewerMock.EXPECT().Output(measurementID1, expectedOpts).Times(1).Return(nil)
+
+	timeMock := mocks.NewMockTime(ctrl)
+	timeMock.EXPECT().Now().Return(defaultCurrentTime).AnyTimes()
+
+	w := new(bytes.Buffer)
+	printer := view.NewPrinter(nil, w, w)
+	ctx := createDefaultContext("http")
+	root := NewRoot(printer, ctx, viewerMock, timeMock, gbMock, nil)
+	os.Args = []string{"globalping", "http", "jsdelivr.com",
+		"from", "Berlin",
+		"--ipv4",
+	}
+	err := root.Cmd.ExecuteContext(context.TODO())
+	assert.NoError(t, err)
+
+	assert.Equal(t, "", w.String())
+
+	expectedCtx := createDefaultExpectedContext("http")
+	expectedCtx.Ipv4 = true
+
+	assert.Equal(t, expectedCtx, ctx)
+}
+
+func Test_Execute_HTTP_IPv6(t *testing.T) {
+	t.Cleanup(sessionCleanup)
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	expectedOpts := createDefaultMeasurementCreate("http")
+	expectedOpts.Options.Protocol = "http"
+	expectedOpts.Options.IPVersion = globalping.IPVersion6
+	expectedOpts.Options.Request = &globalping.RequestOptions{
+		Headers: map[string]string{},
+		Host:    "jsdelivr.com",
+	}
+
+	expectedResponse := createDefaultMeasurementCreateResponse()
+
+	gbMock := mocks.NewMockClient(ctrl)
+	gbMock.EXPECT().CreateMeasurement(expectedOpts).Times(1).Return(expectedResponse, false, nil)
+
+	viewerMock := mocks.NewMockViewer(ctrl)
+	viewerMock.EXPECT().Output(measurementID1, expectedOpts).Times(1).Return(nil)
+
+	timeMock := mocks.NewMockTime(ctrl)
+	timeMock.EXPECT().Now().Return(defaultCurrentTime).AnyTimes()
+
+	w := new(bytes.Buffer)
+	printer := view.NewPrinter(nil, w, w)
+	ctx := createDefaultContext("http")
+	root := NewRoot(printer, ctx, viewerMock, timeMock, gbMock, nil)
+	os.Args = []string{"globalping", "http", "jsdelivr.com",
+		"from", "Berlin",
+		"--ipv6",
+	}
+	err := root.Cmd.ExecuteContext(context.TODO())
+	assert.NoError(t, err)
+
+	assert.Equal(t, "", w.String())
+
+	expectedCtx := createDefaultExpectedContext("http")
+	expectedCtx.Ipv6 = true
+
+	assert.Equal(t, expectedCtx, ctx)
+}
+
 func Test_ParseUrlData(t *testing.T) {
 	urlData, err := parseUrlData("https://cdn.jsdelivr.net:8080/npm/react/?query=3")
 	assert.NoError(t, err)

diff --git a/cmd/mtr.go b/cmd/mtr.go
@@ -84,6 +84,12 @@ func (r *Root) RunMTR(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	if r.ctx.Ipv4 {
+		opts.Options.IPVersion = globalping.IPVersion4
+	} else if r.ctx.Ipv6 {
+		opts.Options.IPVersion = globalping.IPVersion6
+	}
+
 	res, showHelp, err := r.client.CreateMeasurement(opts)
 	if err != nil {
 		if !showHelp {