Merge pull request #38 from shubham-stepsecurity/sm/update

feat(mdm-npm): add gzip compression for stdout/stderr output

Author: ashishkurmi

internal/telemetry/telemetry.go

@@ -2,6 +2,7 @@ package telemetry
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -525,9 +526,18 @@ func uploadToS3(ctx context.Context, log *progress.Logger, payload *Payload) err
 		return fmt.Errorf("marshaling payload: %w", err)
 	}
 
+	// Gzip-compress the payload before upload. The backend signals support by
+	// honoring is_compressed=true on the upload-URL request and appending .gz
+	// to the S3 key, which tells GetTelemetryFromS3 to decompress on read.
+	compressedPayload, err := gzipBytes(payloadJSON)
+	if err != nil {
+		return fmt.Errorf("compressing payload: %w", err)
+	}
+
 	// Request upload URL
-	reqBody, _ := json.Marshal(map[string]string{
-		"device_id": payload.DeviceID,
+	reqBody, _ := json.Marshal(map[string]any{
+		"device_id":     payload.DeviceID,
+		"is_compressed": true,
 	})
 
 	uploadURLEndpoint := fmt.Sprintf("%s/v1/%s/developer-mdm-agent/telemetry/upload-url",
@@ -562,15 +572,15 @@ func uploadToS3(ctx context.Context, log *progress.Logger, payload *Payload) err
 		return fmt.Errorf("empty upload URL in response")
 	}
 
-	// Upload payload to S3 with retry — use a longer timeout since payloads
-	// with npm scan data and execution logs can be several MB.
-	log.Progress("Uploading telemetry to S3 (%d bytes)...", len(payloadJSON))
+	// Upload payload to S3 with retry. Content-Type stays application/json to
+	// match the presigned URL's signed headers — the body is gzipped JSON bytes.
+	log.Progress("Uploading telemetry to S3 (%d bytes)...", len(compressedPayload))
 	s3Client := &http.Client{Timeout: 10 * time.Minute}
 	const maxRetries = 3
 	var putResp *http.Response
 	for attempt := 1; attempt <= maxRetries; attempt++ {
 		uploadStart := time.Now()
-		putReq, reqErr := http.NewRequestWithContext(ctx, http.MethodPut, urlResp.UploadURL, bytes.NewReader(payloadJSON))
+		putReq, reqErr := http.NewRequestWithContext(ctx, http.MethodPut, urlResp.UploadURL, bytes.NewReader(compressedPayload))
 		if reqErr != nil {
 			return fmt.Errorf("creating S3 PUT request: %w", reqErr)
 		}
@@ -598,10 +608,10 @@ func uploadToS3(ctx context.Context, log *progress.Logger, payload *Payload) err
 		if attempt == maxRetries {
 			if err != nil {
 				return fmt.Errorf("uploading to S3 (payload: %d bytes, elapsed: %s, attempts: %d): %w",
-					len(payloadJSON), elapsed, maxRetries, err)
+					len(compressedPayload), elapsed, maxRetries, err)
 			}
 			return fmt.Errorf("S3 upload failed with status %d (payload: %d bytes, attempts: %d)",
-				putResp.StatusCode, len(payloadJSON), maxRetries)
+				putResp.StatusCode, len(compressedPayload), maxRetries)
 		}
 
 		// Log retry and backoff
@@ -654,6 +664,19 @@ func uploadToS3(ctx context.Context, log *progress.Logger, payload *Payload) err
 	return nil
 }
 
+// gzipBytes returns a gzip-compressed copy of the input bytes.
+func gzipBytes(data []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	gz := gzip.NewWriter(&buf)
+	if _, err := gz.Write(data); err != nil {
+		return nil, err
+	}
+	if err := gz.Close(); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
 func resolveSearchDirs(exec executor.Executor, dirs []string) []string {
 	resolved := make([]string, 0, len(dirs))
 	for _, d := range dirs {

internal/telemetry/telemetry_test.go

@@ -0,0 +1,154 @@
+package telemetry
+
+import (
+	"bytes"
+	"compress/gzip"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/step-security/dev-machine-guard/internal/config"
+	"github.com/step-security/dev-machine-guard/internal/progress"
+)
+
+func TestGzipBytes_RoundTrip(t *testing.T) {
+	original := []byte(`{"customer_id":"acme","node_projects":[{"project_path":"/x"}]}`)
+	compressed, err := gzipBytes(original)
+	if err != nil {
+		t.Fatalf("gzipBytes failed: %v", err)
+	}
+	if len(compressed) < 2 || compressed[0] != 0x1f || compressed[1] != 0x8b {
+		t.Fatal("expected gzip magic bytes")
+	}
+
+	gz, err := gzip.NewReader(bytes.NewReader(compressed))
+	if err != nil {
+		t.Fatalf("gzip.NewReader failed: %v", err)
+	}
+	defer gz.Close()
+	got, err := io.ReadAll(gz)
+	if err != nil {
+		t.Fatalf("decompression failed: %v", err)
+	}
+	if !bytes.Equal(got, original) {
+		t.Errorf("round-trip mismatch: got %q, want %q", got, original)
+	}
+}
+
+func TestUploadToS3_SendsCompressedBodyAndIsCompressedFlag(t *testing.T) {
+	var (
+		mu             sync.Mutex
+		uploadURLBody  []byte
+		putBody        []byte
+		putContentType string
+		notifyBody     []byte
+	)
+
+	// Mock S3 PUT endpoint — captures the body the agent uploads.
+	s3Server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, _ := io.ReadAll(r.Body)
+		mu.Lock()
+		putBody = body
+		putContentType = r.Header.Get("Content-Type")
+		mu.Unlock()
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer s3Server.Close()
+
+	// Mock backend — handles upload-URL and process-uploaded calls.
+	backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case strings.HasSuffix(r.URL.Path, "/telemetry/upload-url"):
+			body, _ := io.ReadAll(r.Body)
+			mu.Lock()
+			uploadURLBody = body
+			mu.Unlock()
+			_ = json.NewEncoder(w).Encode(map[string]string{
+				"upload_url": s3Server.URL + "/put",
+				"s3_key":     "developer-mdm/test-customer/dev-1/123.json.gz",
+			})
+		case strings.HasSuffix(r.URL.Path, "/telemetry/process-uploaded"):
+			body, _ := io.ReadAll(r.Body)
+			mu.Lock()
+			notifyBody = body
+			mu.Unlock()
+			w.WriteHeader(http.StatusOK)
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	defer backendServer.Close()
+
+	// Override config globals for the duration of the test.
+	origEndpoint, origCustomer, origKey := config.APIEndpoint, config.CustomerID, config.APIKey
+	config.APIEndpoint = backendServer.URL
+	config.CustomerID = "test-customer"
+	config.APIKey = "test-key"
+	defer func() {
+		config.APIEndpoint, config.CustomerID, config.APIKey = origEndpoint, origCustomer, origKey
+	}()
+
+	payload := &Payload{
+		CustomerID: "test-customer",
+		DeviceID:   "dev-1",
+	}
+
+	if err := uploadToS3(context.Background(), progress.NewLogger(progress.LevelInfo), payload); err != nil {
+		t.Fatalf("uploadToS3 failed: %v", err)
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	// Upload-URL request body must include is_compressed: true.
+	var uploadReq map[string]any
+	if err := json.Unmarshal(uploadURLBody, &uploadReq); err != nil {
+		t.Fatalf("failed to parse upload-URL request body: %v", err)
+	}
+	if uploadReq["device_id"] != "dev-1" {
+		t.Errorf("expected device_id=dev-1, got %v", uploadReq["device_id"])
+	}
+	if uploadReq["is_compressed"] != true {
+		t.Errorf("expected is_compressed=true, got %v", uploadReq["is_compressed"])
+	}
+
+	// PUT body must be gzip-compressed.
+	if len(putBody) < 2 || putBody[0] != 0x1f || putBody[1] != 0x8b {
+		t.Fatalf("expected gzip-compressed PUT body (got %d bytes)", len(putBody))
+	}
+	if putContentType != "application/json" {
+		t.Errorf("expected Content-Type application/json (matches presigned URL), got %q", putContentType)
+	}
+
+	// Decompressing the PUT body should yield the original JSON payload.
+	gz, err := gzip.NewReader(bytes.NewReader(putBody))
+	if err != nil {
+		t.Fatalf("PUT body is not valid gzip: %v", err)
+	}
+	defer gz.Close()
+	decompressed, err := io.ReadAll(gz)
+	if err != nil {
+		t.Fatalf("failed to decompress PUT body: %v", err)
+	}
+	var roundTrip Payload
+	if err := json.Unmarshal(decompressed, &roundTrip); err != nil {
+		t.Fatalf("decompressed body is not valid JSON: %v", err)
+	}
+	if roundTrip.DeviceID != "dev-1" {
+		t.Errorf("decompressed payload device_id mismatch: got %q", roundTrip.DeviceID)
+	}
+
+	// Notify-backend was called with the s3_key returned from the upload-URL endpoint.
+	var notify map[string]string
+	if err := json.Unmarshal(notifyBody, &notify); err != nil {
+		t.Fatalf("failed to parse notify body: %v", err)
+	}
+	if !strings.HasSuffix(notify["s3_key"], ".json.gz") {
+		t.Errorf("expected s3_key with .json.gz suffix, got %q", notify["s3_key"])
+	}
+}