fix: consolidate llm client logic so all requests have consistent headers [IDE-1350] (#113)

Author: andrewrobinsonhodges-snyk

README.md

@@ -86,7 +86,8 @@ Implement the `config.Config` interface to configure the Snyk Code API client fr
 
 Use the Code Scanner to trigger a scan for a Snyk Code workspace using the Bundle Manager.
 
-The Code Scanner exposes a `UploadAndAnalyze` function, which can be used like this:
+The Code Scanner exposes two scanning functions: `UploadAndAnalyze` (which supports Code Consistent Ignores) and 
+`UploadAndAnalyzeLegacy`. These functions may be used like this:
 
 ```go
 import (
@@ -102,7 +103,12 @@ codeScanner := codeClient.NewCodeScanner(
     codeClientHTTP.WithInstrumentor(instrumentor),
     codeClientHTTP.WithErrorReporter(errorReporter),
 )
-codeScanner.UploadAndAnalyze(context.Background(), requestId, target, channelForWalkingFiles, changedFiles)
+if useCodeConsistentIgnores() {
+    codeScanner.UploadAndAnalyze(context.Background(), requestId, target, channelForWalkingFiles, changedFiles)
+} else {
+    codeScanner.UploadAndAnalyzeLegacy(context.Background(), requestId, target, shardKey, files, changedFiles, statusChannel)
+}
+
 ```
 
 

http/http.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/rs/zerolog"
 
+	"github.com/snyk/code-client-go/internal/util/encoding"
 	"github.com/snyk/code-client-go/observability"
 )
 
@@ -100,6 +101,8 @@ var retryErrorCodes = map[int]bool{
 	http.StatusInternalServerError: true,
 }
 
+const NoRequestId = ""
+
 func (s *httpClient) Do(req *http.Request) (*http.Response, error) {
 	span := s.instrumentor.StartSpan(req.Context(), "http.Do")
 	defer s.instrumentor.Finish(span)
@@ -160,14 +163,43 @@ func NewDefaultClientFactory() HTTPClientFactory {
 	return clientFunc
 }
 
-func AddDefaultHeaders(req *http.Request, requestId string, orgId string) {
+func AddDefaultHeaders(req *http.Request, requestId string, orgId string, method string) {
 	// if requestId is empty it will be enriched from the Gateway
 	if len(requestId) > 0 {
 		req.Header.Set("snyk-request-id", requestId)
 	}
 	if len(orgId) > 0 {
 		req.Header.Set("snyk-org-name", orgId)
 	}
+
+	// https://www.keycdn.com/blog/http-cache-headers
 	req.Header.Set("Cache-Control", "private, max-age=0, no-cache")
-	req.Header.Set("Content-Type", "application/json")
+
+	if mustBeEncoded(method) {
+		req.Header.Set("Content-Type", "application/octet-stream")
+		req.Header.Set("Content-Encoding", "gzip")
+	} else {
+		req.Header.Set("Content-Type", "application/json")
+	}
+}
+
+// EncodeIfNeeded returns a byte buffer for the requestBody. Depending on the request method, it may encode the buffer.
+// (See http.mustBeEncoded for the list of methods which require encoding the request body.)
+func EncodeIfNeeded(method string, requestBody []byte) (*bytes.Buffer, error) {
+	b := new(bytes.Buffer)
+	if mustBeEncoded(method) {
+		enc := encoding.NewEncoder(b)
+		_, err := enc.Write(requestBody)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		b = bytes.NewBuffer(requestBody)
+	}
+	return b, nil
+}
+
+// mustBeEncoded returns true if the request method requires the request body to be encoded.
+func mustBeEncoded(method string) bool {
+	return method == http.MethodPost || method == http.MethodPut
 }

internal/analysis/analysis_legacy.go

@@ -28,10 +28,9 @@ import (
 	"net/url"
 	"strings"
 
-	"github.com/snyk/code-client-go/scan"
-
 	codeClientHTTP "github.com/snyk/code-client-go/http"
 	"github.com/snyk/code-client-go/sarif"
+	"github.com/snyk/code-client-go/scan"
 )
 
 // Legacy analysis types and constants
@@ -76,15 +75,21 @@ type FailedError struct {
 func (e FailedError) Error() string { return e.Msg }
 
 // Legacy analysis helper functions
-func (a *analysisOrchestrator) newRequestContext() requestContext {
+func (a *analysisOrchestrator) newRequestContext(ctx context.Context) requestContext {
 	unknown := "unknown"
 	orgId := unknown
 	if a.config.Organization() != "" {
 		orgId = a.config.Organization()
 	}
 
+	initiator := unknown
+	contextInitiator, ok := scan.ScanSourceFromContext(ctx)
+	if ok {
+		initiator = string(contextInitiator)
+	}
+
 	return requestContext{
-		Initiator: "IDE",
+		Initiator: initiator,
 		Flow:      "language-server",
 		Org: requestContextOrg{
 			Name:        unknown,
@@ -94,15 +99,15 @@ func (a *analysisOrchestrator) newRequestContext() requestContext {
 	}
 }
 
-func (a *analysisOrchestrator) createRequestBody(bundleHash, shardKey string, limitToFiles []string, severity int) ([]byte, error) {
+func (a *analysisOrchestrator) createRequestBody(ctx context.Context, bundleHash, shardKey string, limitToFiles []string, severity int) ([]byte, error) {
 	request := Request{
 		Key: RequestKey{
 			Type:         "file",
 			Hash:         bundleHash,
 			LimitToFiles: limitToFiles,
 		},
 		Legacy:          false,
-		AnalysisContext: a.newRequestContext(),
+		AnalysisContext: a.newRequestContext(ctx),
 	}
 	if len(shardKey) > 0 {
 		request.Key.Shard = shardKey
@@ -144,7 +149,7 @@ func (a *analysisOrchestrator) RunLegacyTest(ctx context.Context, bundleHash str
 	a.logger.Debug().Str("method", method).Str("bundleHash", bundleHash).Msg("API: Retrieving analysis for bundle")
 	defer a.logger.Debug().Str("method", method).Str("bundleHash", bundleHash).Msg("API: Retrieving analysis done")
 
-	requestBody, err := a.createRequestBody(bundleHash, shardKey, limitToFiles, severity)
+	requestBody, err := a.createRequestBody(ctx, bundleHash, shardKey, limitToFiles, severity)
 	if err != nil {
 		a.logger.Err(err).Str("method", method).Str("requestBody", string(requestBody)).Msg("error creating request body")
 		return nil, scan.LegacyScanStatus{}, err
@@ -158,12 +163,13 @@ func (a *analysisOrchestrator) RunLegacyTest(ctx context.Context, bundleHash str
 
 	// Create HTTP request
 	analysisUrl := baseUrl + "/analysis"
-	req, err := http.NewRequestWithContext(span.Context(), http.MethodPost, analysisUrl, bytes.NewBuffer(requestBody))
+	httpMethod := http.MethodPost
+	req, err := http.NewRequestWithContext(span.Context(), httpMethod, analysisUrl, bytes.NewBuffer(requestBody))
 	if err != nil {
 		a.logger.Err(err).Str("method", method).Msg("error creating HTTP request")
 		return nil, scan.LegacyScanStatus{}, err
 	}
-	codeClientHTTP.AddDefaultHeaders(req, span.GetTraceId(), a.config.Organization())
+	codeClientHTTP.AddDefaultHeaders(req, span.GetTraceId(), a.config.Organization(), httpMethod)
 
 	// Make HTTP call
 	resp, err := a.httpClient.Do(req)

internal/analysis/analysis_legacy_test.go

@@ -530,7 +530,7 @@ func TestAnalysis_CreateRequestBody(t *testing.T) {
 
 	// run method under test
 	_, _, err := analysisOrchestrator.RunLegacyTest(
-		t.Context(),
+		scan.NewContextWithScanSource(t.Context(), scan.IDE),
 		bundleHash,
 		shardKey,
 		limitToFiles,
@@ -568,7 +568,7 @@ func TestAnalysis_CreateRequestBody(t *testing.T) {
 
 	// Validate analysisContext
 	analysisContext := request["analysisContext"].(map[string]interface{})
-	assert.Equal(t, "IDE", analysisContext["initiator"])
+	assert.Equal(t, string(scan.IDE), analysisContext["initiator"])
 	assert.Equal(t, "language-server", analysisContext["flow"])
 
 	org := analysisContext["org"].(map[string]interface{})

internal/deepcode/client.go

@@ -17,7 +17,6 @@
 package deepcode
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
@@ -30,7 +29,6 @@ import (
 
 	"github.com/rs/zerolog"
 	"github.com/snyk/code-client-go/config"
-	"github.com/snyk/code-client-go/internal/util/encoding"
 
 	codeClientHTTP "github.com/snyk/code-client-go/http"
 	"github.com/snyk/code-client-go/observability"
@@ -217,7 +215,7 @@ func (s *deepcodeClient) Request(
 		return nil, err
 	}
 
-	bodyBuffer, err := s.encodeIfNeeded(method, requestBody)
+	bodyBuffer, err := codeClientHTTP.EncodeIfNeeded(method, requestBody)
 	if err != nil {
 		return nil, err
 	}
@@ -227,7 +225,7 @@ func (s *deepcodeClient) Request(
 		return nil, err
 	}
 
-	s.addHeaders(method, req)
+	codeClientHTTP.AddDefaultHeaders(req, codeClientHTTP.NoRequestId, s.config.Organization(), method)
 
 	response, err := s.httpClient.Do(req)
 	if err != nil {
@@ -255,41 +253,6 @@ func (s *deepcodeClient) Request(
 	return responseBody, nil
 }
 
-func (s *deepcodeClient) addHeaders(method string, req *http.Request) {
-	// Setting a chosen org name for the request
-	org := s.config.Organization()
-	if org != "" {
-		req.Header.Set("snyk-org-name", org)
-	}
-	// https://www.keycdn.com/blog/http-cache-headers
-	req.Header.Set("Cache-Control", "private, max-age=0, no-cache")
-	if s.mustBeEncoded(method) {
-		req.Header.Set("Content-Type", "application/octet-stream")
-		req.Header.Set("Content-Encoding", "gzip")
-	} else {
-		req.Header.Set("Content-Type", "application/json")
-	}
-}
-
-func (s *deepcodeClient) encodeIfNeeded(method string, requestBody []byte) (*bytes.Buffer, error) {
-	b := new(bytes.Buffer)
-	mustBeEncoded := s.mustBeEncoded(method)
-	if mustBeEncoded {
-		enc := encoding.NewEncoder(b)
-		_, err := enc.Write(requestBody)
-		if err != nil {
-			return nil, err
-		}
-	} else {
-		b = bytes.NewBuffer(requestBody)
-	}
-	return b, nil
-}
-
-func (s *deepcodeClient) mustBeEncoded(method string) bool {
-	return method == http.MethodPost || method == http.MethodPut
-}
-
 func (s *deepcodeClient) checkResponseCode(r *http.Response) error {
 	if r.StatusCode >= 200 && r.StatusCode <= 299 {
 		return nil

llm/api_client.go

@@ -1,7 +1,6 @@
 package llm
 
 import (
-	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -12,7 +11,7 @@ import (
 	"net/url"
 	"strings"
 
-	codeClientHTTP "github.com/snyk/code-client-go/http"
+	http2 "github.com/snyk/code-client-go/http"
 )
 
 var (
@@ -70,13 +69,20 @@ func (d *DeepCodeLLMBindingImpl) submitRequest(ctx context.Context, url *url.URL
 	span := d.instrumentor.StartSpan(ctx, "code.SubmitRequest")
 	defer span.Finish()
 
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url.String(), bytes.NewBuffer(requestBody))
+	// Encode the request body
+	bodyBuffer, err := http2.EncodeIfNeeded(http.MethodPost, requestBody)
+	if err != nil {
+		logger.Err(err).Str("requestBody", string(requestBody)).Msg("error encoding request body")
+		return nil, err
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url.String(), bodyBuffer)
 	if err != nil {
 		logger.Err(err).Str("requestBody", string(requestBody)).Msg("error creating request")
 		return nil, err
 	}
 
-	codeClientHTTP.AddDefaultHeaders(req, span.GetTraceId(), orgId)
+	http2.AddDefaultHeaders(req, http2.NoRequestId, orgId, http.MethodPost)
 
 	resp, err := d.httpClientFunc().Do(req) //nolint:bodyclose // this seems to be a false positive
 	if err != nil {

llm/api_client_test.go

@@ -10,10 +10,10 @@ import (
 	"testing"
 
 	"github.com/rs/zerolog"
+	http2 "github.com/snyk/code-client-go/http"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	codeClientHTTP "github.com/snyk/code-client-go/http"
 	"github.com/snyk/code-client-go/observability"
 )
 
@@ -264,7 +264,7 @@ func testLogger(t *testing.T) *zerolog.Logger {
 func TestAddDefaultHeadersWithExistingHeaders(t *testing.T) {
 	req := &http.Request{Header: http.Header{"Existing-Header": {"existing-value"}}}
 
-	codeClientHTTP.AddDefaultHeaders(req, "", "")
+	http2.AddDefaultHeaders(req, http2.NoRequestId, "", http.MethodGet)
 
 	cacheControl := req.Header.Get("Cache-Control")
 	contentType := req.Header.Get("Content-Type")

scan.go

@@ -66,6 +66,8 @@ type CodeScanner interface {
 		changedFiles map[string]bool,
 	) (*sarif.SarifResponse, string, error)
 
+	// UploadAndAnalyzeLegacy runs the legacy scanner (no consistent ignores)
+	// ctx may include a scan.ScanSource value for use in the requestContext (see analysis_legacy.go)
 	UploadAndAnalyzeLegacy(
 		ctx context.Context,
 		requestId string,

scan/resultmetadata.go

@@ -0,0 +1,33 @@
+package scan
+
+import "context"
+
+type ResultMetaData struct {
+	FindingsUrl string
+	WebUiUrl    string
+}
+
+type ScanSource string
+
+func (s ScanSource) String() string {
+	return string(s)
+}
+
+const (
+	LLM ScanSource = "LLM"
+	IDE ScanSource = "IDE"
+	CLI ScanSource = "CLI"
+)
+
+type scanSourceKeyType int
+
+var scanSourceKey scanSourceKeyType
+
+func NewContextWithScanSource(ctx context.Context, source ScanSource) context.Context {
+	return context.WithValue(ctx, scanSourceKey, source)
+}
+
+func ScanSourceFromContext(ctx context.Context) (ScanSource, bool) {
+	s, ok := ctx.Value(scanSourceKey).(ScanSource)
+	return s, ok
+}