CLI | Add redirected request mechanism (AST-37864) (#712)

* fix redirect mechanism * check * check * add const error variables --------- Co-authored-by: AlvoBen <[email protected]>
Checkmarx · Apr 22, 2024 · e446185 · e446185
1 parent 4350672
commit e446185
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 12 deletions.
diff --git a/go.mod b/go.mod
@@ -15,6 +15,7 @@ require (
 	github.com/spf13/cobra v1.8.0
 	github.com/spf13/viper v1.18.2
 	github.com/stretchr/testify v1.8.4
+	golang.org/x/net v0.23.0
 	github.com/tomnomnom/linkheader v0.0.0-20180905144013-02ca5825eb80
 	golang.org/x/crypto v0.22.0
 	golang.org/x/text v0.14.0

diff --git a/internal/errors/application-errors.go b/internal/errors/application-errors.go
@@ -6,4 +6,6 @@ const (
 
 const (
 	FailedToGetApplication = "Failed to get application"
+	RedirectURLNotFound    = "redirect URL not found in response"
+	HTTPMethodNotFound     = "HTTP method not found in request"
 )
diff --git a/internal/wrappers/client.go b/internal/wrappers/client.go
@@ -14,6 +14,7 @@ import (
 	"strings"
 	"time"
 
+	applicationErrors "github.com/checkmarx/ast-cli/internal/errors"
 	"github.com/golang-jwt/jwt"
 
 	"github.com/checkmarx/ast-cli/internal/logger"
@@ -580,18 +581,7 @@ func request(client *http.Client, req *http.Request, responseBody bool) (*http.R
 		}
 		if resp != nil && err == nil {
 			if hasRedirectStatusCode(resp) {
-				redirectURL := resp.Header.Get("Location")
-				if redirectURL == "" {
-					return nil, fmt.Errorf("redirect URL not found in response")
-				}
-				method := GetHTTPMethod(req)
-				if method == "" {
-					return nil, fmt.Errorf("method not found in request")
-				}
-				req, err = http.NewRequest(method, redirectURL, bytes.NewReader(body))
-				if err != nil {
-					return nil, err
-				}
+				req, err = handleRedirect(resp, req, body)
 				continue
 			}
 			logger.PrintResponse(resp, responseBody)
@@ -603,6 +593,40 @@ func request(client *http.Client, req *http.Request, responseBody bool) (*http.R
 	return nil, err
 }
 
+func handleRedirect(resp *http.Response, req *http.Request, body []byte) (*http.Request, error) {
+	redirectURL := resp.Header.Get("Location")
+	if redirectURL == "" {
+		return nil, fmt.Errorf(applicationErrors.RedirectURLNotFound)
+	}
+
+	method := GetHTTPMethod(req)
+	if method == "" {
+		return nil, fmt.Errorf(applicationErrors.HTTPMethodNotFound)
+	}
+
+	newReq, err := recreateRequest(req, method, redirectURL, body)
+	if err != nil {
+		return nil, err
+	}
+
+	return newReq, nil
+}
+
+func recreateRequest(oldReq *http.Request, method, redirectURL string, body []byte) (*http.Request, error) {
+	newReq, err := http.NewRequest(method, redirectURL, io.NopCloser(bytes.NewBuffer(body)))
+	if err != nil {
+		return nil, err
+	}
+
+	for key, values := range oldReq.Header {
+		for _, value := range values {
+			newReq.Header.Add(key, value)
+		}
+	}
+
+	return newReq, nil
+}
+
 func GetHTTPMethod(req *http.Request) string {
 	switch req.Method {
 	case http.MethodGet: