311 lines
10 KiB
Go
311 lines
10 KiB
Go
package distributor
|
|
|
|
import (
|
|
"archive/tar"
|
|
"compress/gzip"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/weatherreporter/internal/config"
|
|
)
|
|
|
|
const oversizedRemoteDiagnostic = "REMOTE-DIAGNOSTIC"
|
|
|
|
func TestUploadUsesProductionHTTPBoundary(t *testing.T) {
|
|
const token = "test-upload-token"
|
|
var uploadCalls, statusCalls int
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch {
|
|
case r.Method == http.MethodPost && r.URL.Path == "/prefix/v1/pipelines/weather/upload":
|
|
uploadCalls++
|
|
if got := r.Header.Get("Authorization"); got != "Bearer "+token {
|
|
t.Fatalf("authorization = %q", got)
|
|
}
|
|
if got := r.Header.Get("Idempotency-Key"); got != "bundle-key" {
|
|
t.Fatalf("idempotency key = %q", got)
|
|
}
|
|
if got := r.Header.Get("Content-Type"); got != "application/gzip" {
|
|
t.Fatalf("content type = %q", got)
|
|
}
|
|
verifyUploadedArchive(t, r.Body, "daily/report.md", "report body")
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusAccepted)
|
|
_, _ = io.WriteString(w, `{"run_id":"run-123","status":"accepted"}`)
|
|
case r.Method == http.MethodGet && r.URL.Path == "/prefix/runs/run-123":
|
|
statusCalls++
|
|
if got := r.Header.Get("Authorization"); got != "Bearer "+token {
|
|
t.Fatalf("authorization = %q", got)
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"run_id":"run-123","pipeline_id":"weather","status":"succeeded","report":{"detail":"REMOTE-DETAIL"}}`)
|
|
default:
|
|
t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
client := productionClient(t, server.URL+"/prefix", token)
|
|
result, err := client.Upload(context.Background(), productionUploadRequest(t))
|
|
if err != nil || uploadCalls != 1 || statusCalls != 1 || result.RunID != "run-123" || result.Status != "succeeded" || result.UploadStatus != "accepted" || result.RunStatus == nil || result.RunStatus.PipelineID != "weather" || len(result.RunStatus.Report) != 0 {
|
|
t.Fatalf("result/error/calls = %#v/%v/%d/%d", result, err, uploadCalls, statusCalls)
|
|
}
|
|
}
|
|
|
|
func TestUploadClassifiesRemoteHTTPDiagnostics(t *testing.T) {
|
|
const token = "test-upload-token"
|
|
const remote = oversizedRemoteDiagnostic
|
|
for _, tt := range []struct {
|
|
name string
|
|
handle func(http.ResponseWriter, *http.Request)
|
|
check func(t *testing.T, result UploadResult, err error)
|
|
}{
|
|
{
|
|
name: "upload failure",
|
|
handle: func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
t.Fatalf("method = %s", r.Method)
|
|
}
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
_, _ = io.WriteString(w, `{"error":"REMOTE-DIAGNOSTIC","retryable":true}`)
|
|
},
|
|
check: func(t *testing.T, _ UploadResult, err error) {
|
|
t.Helper()
|
|
var remoteErr *RemoteResponseError
|
|
if err == nil || !errors.As(err, &remoteErr) || remoteErr.StatusCode != http.StatusBadRequest || !remoteErr.Retryable {
|
|
t.Fatalf("error = %T %v", err, err)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "status failure",
|
|
handle: func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method == http.MethodPost {
|
|
w.WriteHeader(http.StatusAccepted)
|
|
_, _ = io.WriteString(w, `{"run_id":"run-123","status":"accepted"}`)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
_, _ = io.WriteString(w, remote)
|
|
},
|
|
check: func(t *testing.T, result UploadResult, err error) {
|
|
t.Helper()
|
|
if err != nil || result.Status != "accepted" || result.StatusError != "distributor request failed with HTTP status 500" {
|
|
t.Fatalf("result/error = %#v/%v", result, err)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "failed run",
|
|
handle: func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method == http.MethodPost {
|
|
w.WriteHeader(http.StatusAccepted)
|
|
_, _ = io.WriteString(w, `{"run_id":"run-123","status":"accepted"}`)
|
|
return
|
|
}
|
|
_, _ = io.WriteString(w, `{"run_id":"run-123","status":"failed","error":"REMOTE-DIAGNOSTIC","report":{"detail":"REMOTE-DIAGNOSTIC"}}`)
|
|
},
|
|
check: func(t *testing.T, result UploadResult, err error) {
|
|
t.Helper()
|
|
if err == nil || result.Status != "failed" || result.RunStatus == nil || result.RunStatus.Error != "distributor reported a failed run" || len(result.RunStatus.Report) != 0 {
|
|
t.Fatalf("result/error = %#v/%v", result, err)
|
|
}
|
|
},
|
|
},
|
|
} {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(tt.handle))
|
|
defer server.Close()
|
|
result, err := productionClient(t, server.URL, token).Upload(context.Background(), productionUploadRequest(t))
|
|
tt.check(t, result, err)
|
|
for _, value := range []string{fmt.Sprint(result), fmt.Sprint(err)} {
|
|
if strings.Contains(value, remote) || strings.Contains(value, token) {
|
|
t.Fatalf("normal diagnostic leaked remote value: %q", value)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestUploadBoundsHTTPResponses(t *testing.T) {
|
|
for _, tt := range []struct {
|
|
name string
|
|
response func(size int) string
|
|
statusCode int
|
|
statusBody func(size int) string
|
|
check func(t *testing.T, result UploadResult, err error, overflow bool)
|
|
}{
|
|
{
|
|
name: "accepted response",
|
|
response: func(size int) string {
|
|
return paddedJSON(t, `{"run_id":"run-123","status":"accepted","detail":"REMOTE-DIAGNOSTIC"}`, size)
|
|
},
|
|
statusBody: func(_ int) string {
|
|
return `{"run_id":"run-123","status":"succeeded"}`
|
|
},
|
|
check: func(t *testing.T, result UploadResult, err error, overflow bool) {
|
|
t.Helper()
|
|
if overflow {
|
|
if !errors.Is(err, errDistributorResponseTooLarge) || result.RunID != "" {
|
|
t.Fatalf("overflow result/error = %#v/%v", result, err)
|
|
}
|
|
return
|
|
}
|
|
if err != nil || result.Status != "succeeded" {
|
|
t.Fatalf("bounded result/error = %#v/%v", result, err)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "status report",
|
|
response: func(_ int) string {
|
|
return `{"run_id":"run-123","status":"accepted"}`
|
|
},
|
|
statusBody: func(size int) string { return statusReportBody(t, size) },
|
|
check: func(t *testing.T, result UploadResult, err error, overflow bool) {
|
|
t.Helper()
|
|
if overflow {
|
|
if err != nil || result.Status != "accepted" || result.StatusError != errDistributorResponseTooLarge.Error() {
|
|
t.Fatalf("overflow result/error = %#v/%v", result, err)
|
|
}
|
|
return
|
|
}
|
|
if err != nil || result.Status != "succeeded" || result.RunStatus == nil || len(result.RunStatus.Report) != 0 {
|
|
t.Fatalf("bounded result/error = %#v/%v", result, err)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "error response",
|
|
response: func(size int) string { return repeatedToLength(oversizedRemoteDiagnostic, size) },
|
|
statusCode: http.StatusBadRequest,
|
|
check: func(t *testing.T, result UploadResult, err error, overflow bool) {
|
|
t.Helper()
|
|
if overflow {
|
|
if !errors.Is(err, errDistributorResponseTooLarge) || result.RunID != "" {
|
|
t.Fatalf("overflow result/error = %#v/%v", result, err)
|
|
}
|
|
return
|
|
}
|
|
var remoteErr *RemoteResponseError
|
|
if !errors.As(err, &remoteErr) || remoteErr.StatusCode != http.StatusBadRequest {
|
|
t.Fatalf("bounded result/error = %#v/%v", result, err)
|
|
}
|
|
},
|
|
},
|
|
} {
|
|
for _, overflow := range []bool{false, true} {
|
|
t.Run(tt.name+"/"+map[bool]string{false: "limit", true: "over-limit"}[overflow], func(t *testing.T) {
|
|
size := int(maxDistributorResponseBytes)
|
|
if overflow {
|
|
size++
|
|
}
|
|
var uploadCalls int
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method == http.MethodPost {
|
|
uploadCalls++
|
|
statusCode := tt.statusCode
|
|
if statusCode == 0 {
|
|
statusCode = http.StatusAccepted
|
|
}
|
|
w.WriteHeader(statusCode)
|
|
_, _ = io.WriteString(w, tt.response(size))
|
|
return
|
|
}
|
|
_, _ = io.WriteString(w, tt.statusBody(size))
|
|
}))
|
|
defer server.Close()
|
|
|
|
result, err := productionClient(t, server.URL, "test-upload-token").Upload(context.Background(), productionUploadRequest(t))
|
|
tt.check(t, result, err, overflow)
|
|
if strings.Contains(fmt.Sprint(result), oversizedRemoteDiagnostic) || strings.Contains(fmt.Sprint(err), oversizedRemoteDiagnostic) {
|
|
t.Fatalf("result/error leaked oversized response detail: %#v/%v", result, err)
|
|
}
|
|
if uploadCalls != 1 {
|
|
t.Fatalf("upload calls = %d, want one", uploadCalls)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func productionClient(t *testing.T, endpoint, token string) *Client {
|
|
t.Helper()
|
|
cfg := config.Defaults().Notify.Distributor
|
|
cfg.Endpoint = endpoint
|
|
cfg.Timeout = 0
|
|
t.Setenv(cfg.TokenEnv, token)
|
|
return New(cfg)
|
|
}
|
|
|
|
func productionUploadRequest(t *testing.T) UploadRequest {
|
|
t.Helper()
|
|
path := filepath.Join(t.TempDir(), "report.md")
|
|
if err := os.WriteFile(path, []byte("report body"), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return UploadRequest{
|
|
PipelineID: "weather", BundleID: "bundle", IdempotencyKey: "bundle-key",
|
|
Files: []UploadFile{{SourcePath: path, BundlePath: "daily/report.md"}},
|
|
CreatedAt: time.Date(2026, 6, 7, 12, 0, 0, 0, time.UTC),
|
|
}
|
|
}
|
|
|
|
func verifyUploadedArchive(t *testing.T, body io.Reader, wantPath, wantContents string) {
|
|
t.Helper()
|
|
reader, err := gzip.NewReader(body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer reader.Close()
|
|
archive := tar.NewReader(reader)
|
|
for {
|
|
header, err := archive.Next()
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
}
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if header.Name != wantPath {
|
|
continue
|
|
}
|
|
contents, err := io.ReadAll(archive)
|
|
if err != nil || string(contents) != wantContents {
|
|
t.Fatalf("archive file contents/error = %q/%v", contents, err)
|
|
}
|
|
return
|
|
}
|
|
t.Fatalf("archive did not contain %q", wantPath)
|
|
}
|
|
|
|
func paddedJSON(t *testing.T, value string, size int) string {
|
|
t.Helper()
|
|
if len(value) > size {
|
|
t.Fatalf("JSON length = %d, exceeds requested size %d", len(value), size)
|
|
}
|
|
return value + strings.Repeat(" ", size-len(value))
|
|
}
|
|
|
|
func statusReportBody(t *testing.T, size int) string {
|
|
t.Helper()
|
|
const prefix = `{"run_id":"run-123","pipeline_id":"weather","status":"succeeded","report":"`
|
|
const suffix = `"}`
|
|
if len(prefix)+len(suffix) > size {
|
|
t.Fatalf("status response exceeds requested size %d", size)
|
|
}
|
|
return prefix + repeatedToLength(oversizedRemoteDiagnostic, size-len(prefix)-len(suffix)) + suffix
|
|
}
|
|
|
|
func repeatedToLength(value string, size int) string {
|
|
return strings.Repeat(value, size/len(value)+1)[:size]
|
|
}
|